| 45 | |
| 46 | |
| 47 | class Encoder(nn.Module): |
| 48 | |
| 49 | def __init__(self, in_dim, emb_dim, num_layer, kernel='gcn', drop_ratio=0, |
| 50 | act='relu', norm='batchnorm', concat=True, last_act=True,aggr='mean'): |
| 51 | super().__init__() |
| 52 | |
| 53 | self.num_layer = num_layer |
| 54 | self.emb_dim = [in_dim] + [emb_dim] * num_layer |
| 55 | # just try |
| 56 | # self.emb_dim = [in_dim] + [emb_dim] * (num_layer - 1) + [768] |
| 57 | |
| 58 | self.drop_ratio = drop_ratio |
| 59 | self.norm = norm |
| 60 | self.concat = concat |
| 61 | |
| 62 | self.encs = torch.nn.ModuleList() |
| 63 | self.acts = torch.nn.ModuleList() |
| 64 | if self.norm: |
| 65 | self.norms = torch.nn.ModuleList() |
| 66 | |
| 67 | for i in range(self.num_layer): |
| 68 | if kernel == 'gcn': |
| 69 | conv = GCNConv(self.emb_dim[i], self.emb_dim[i + 1], normalize=True, add_self_loops=True) |
| 70 | elif kernel == 'gin': |
| 71 | conv = GINConv(LinearPred(self.emb_dim[i], self.emb_dim[i + 1], self.emb_dim[i + 1], 1)) |
| 72 | elif kernel == 'gin2': |
| 73 | conv = GINConv(LinearPred(self.emb_dim[i], self.emb_dim[i + 1], self.emb_dim[i + 1], 2)) |
| 74 | self.encs.append(conv) |
| 75 | if i == self.num_layer - 1 and not last_act: |
| 76 | act = None |
| 77 | self.acts.append(obtain_act(act)) |
| 78 | |
| 79 | if self.norm: |
| 80 | self.norms.append(obtain_norm(self.norm)(self.emb_dim[i + 1])) |
| 81 | |
| 82 | def forward(self, x, edge_index, edge_weight=None, batch=None): |
| 83 | |
| 84 | xs = [] |
| 85 | for i in range(self.num_layer): |
| 86 | x = self.encs[i](x, edge_index, edge_weight) |
| 87 | x = self.norms[i](x) if self.norm else x |
| 88 | x = F.dropout(self.acts[i](x), self.drop_ratio, training=self.training) |
| 89 | xs.append(x) |
| 90 | |
| 91 | if batch is not None: |
| 92 | xs = [global_add_pool(x, batch) for x in xs] |
| 93 | |
| 94 | if self.concat: |
| 95 | x = torch.concat(xs, dim=1) |
| 96 | else: |
| 97 | x = xs[-1] |
| 98 | |
| 99 | return x |
| 100 | |
| 101 | |
| 102 | class GraphCL(nn.Module): |