| 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 | |