Node classification network for graphs
| 25 | |
| 26 | |
| 27 | class NodeClassifier(nn.Module): |
| 28 | """ |
| 29 | Node classification network for graphs |
| 30 | """ |
| 31 | def __init__(self, embedding, num_class, normalization=False): |
| 32 | super(NodeClassifier, self).__init__() |
| 33 | if normalization: |
| 34 | embedding = embedding / np.linalg.norm(embedding, axis=1, keepdims=True) |
| 35 | embedding = torch.as_tensor(embedding) |
| 36 | self.embeddings = nn.Embedding.from_pretrained(embedding, freeze=True) |
| 37 | self.linear = nn.Linear(embedding.size(1), num_class, bias=True) |
| 38 | |
| 39 | def forward(self, indexes): |
| 40 | x = self.embeddings(indexes) |
| 41 | x = self.linear(x) |
| 42 | return x |
| 43 | |
| 44 | |
| 45 | class LinkPredictor(nn.Module): |
no outgoing calls
no test coverage detected