| 10 | from torch_scatter import scatter_mean |
| 11 | |
| 12 | class GNN(torch.nn.Module): |
| 13 | |
| 14 | def __init__(self, num_vocab, max_seq_len, node_encoder, num_layer = 5, emb_dim = 300, |
| 15 | gnn_type = 'gin', virtual_node = True, residual = False, drop_ratio = 0.5, JK = "last", graph_pooling = "mean", num_class=0, edge_attr_dim=2): |
| 16 | ''' |
| 17 | num_tasks (int): number of labels to be predicted |
| 18 | virtual_node (bool): whether to add virtual node or not |
| 19 | ''' |
| 20 | |
| 21 | super(GNN, self).__init__() |
| 22 | |
| 23 | self.num_class = num_class # if we do classification |
| 24 | self.num_layer = num_layer |
| 25 | self.drop_ratio = drop_ratio |
| 26 | self.JK = JK |
| 27 | self.emb_dim = emb_dim |
| 28 | self.num_vocab = num_vocab |
| 29 | self.max_seq_len = max_seq_len |
| 30 | self.graph_pooling = graph_pooling |
| 31 | |
| 32 | if self.num_layer < 2: |
| 33 | raise ValueError("Number of GNN layers must be greater than 1.") |
| 34 | |
| 35 | ### GNN to generate node embeddings |
| 36 | if virtual_node: |
| 37 | self.gnn_node = GNN_node_Virtualnode(num_layer, emb_dim, node_encoder, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type, edge_attr_dim=edge_attr_dim) |
| 38 | else: |
| 39 | self.gnn_node = GNN_node(num_layer, emb_dim, node_encoder, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type, edge_attr_dim=edge_attr_dim) |
| 40 | |
| 41 | |
| 42 | ### Pooling function to generate whole-graph embeddings |
| 43 | if self.graph_pooling == "sum": |
| 44 | self.pool = global_add_pool |
| 45 | elif self.graph_pooling == "mean": |
| 46 | self.pool = global_mean_pool |
| 47 | elif self.graph_pooling == "max": |
| 48 | self.pool = global_max_pool |
| 49 | elif self.graph_pooling == "attention": |
| 50 | self.pool = GlobalAttention(gate_nn = torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim), torch.nn.BatchNorm1d(2*emb_dim), torch.nn.ReLU(), torch.nn.Linear(2*emb_dim, 1))) |
| 51 | elif self.graph_pooling == "set2set": |
| 52 | self.pool = Set2Set(emb_dim, processing_steps = 2) |
| 53 | else: |
| 54 | raise ValueError("Invalid graph pooling type.") |
| 55 | |
| 56 | self.graph_pred_linear_list = torch.nn.ModuleList() |
| 57 | |
| 58 | if self.num_class > 0: |
| 59 | if graph_pooling == "set2set": |
| 60 | self.graph_pred_linear = torch.nn.Linear(2*self.emb_dim, self.num_class) |
| 61 | else: |
| 62 | self.graph_pred_linear = torch.nn.Linear(self.emb_dim, self.num_class) |
| 63 | else: |
| 64 | if graph_pooling == "set2set": |
| 65 | for i in range(max_seq_len): |
| 66 | self.graph_pred_linear_list.append(torch.nn.Linear(2*emb_dim, self.num_vocab)) |
| 67 | |
| 68 | else: |
| 69 | if self.num_vocab == 1: |