Link prediction network for graphs / knowledge graphs
| 43 | |
| 44 | |
| 45 | class LinkPredictor(nn.Module): |
| 46 | """ |
| 47 | Link prediction network for graphs / knowledge graphs |
| 48 | """ |
| 49 | def __init__(self, score_function, *embeddings, **kwargs): |
| 50 | super(LinkPredictor, self).__init__() |
| 51 | if isinstance(score_function, types.FunctionType): |
| 52 | self.score_function = score_function |
| 53 | else: |
| 54 | self.score_function = getattr(LinkPredictor, score_function) |
| 55 | self.kwargs = kwargs |
| 56 | self.embeddings = nn.ModuleList() |
| 57 | for embedding in embeddings: |
| 58 | embedding = torch.as_tensor(embedding) |
| 59 | embedding = nn.Embedding.from_pretrained(embedding, freeze=True) |
| 60 | self.embeddings.append(embedding) |
| 61 | |
| 62 | def forward(self, *indexes): |
| 63 | assert len(indexes) == len(self.embeddings) |
| 64 | vectors = [] |
| 65 | for index, embedding in zip(indexes, self.embeddings): |
| 66 | vectors.append(embedding(index)) |
| 67 | return self.score_function(*vectors, **self.kwargs) |
| 68 | |
| 69 | @staticmethod |
| 70 | def LINE(heads, tails): |
| 71 | x = heads * tails |
| 72 | score = x.sum(dim=1) |
| 73 | return score |
| 74 | |
| 75 | DeepWalk = LINE |
| 76 | |
| 77 | @staticmethod |
| 78 | def TransE(heads, relations, tails, margin=12): |
| 79 | x = heads + relations - tails |
| 80 | score = margin - x.norm(p=1, dim=1) |
| 81 | return score |
| 82 | |
| 83 | @staticmethod |
| 84 | def RotatE(heads, relations, tails, margin=12): |
| 85 | dim = heads.size(1) // 2 |
| 86 | |
| 87 | head_re, head_im = heads.view(-1, dim, 2).permute(2, 0, 1) |
| 88 | tail_re, tail_im = tails.view(-1, dim, 2).permute(2, 0, 1) |
| 89 | relations = relations[:, :dim] |
| 90 | relation_re, relation_im = torch.cos(relations), torch.sin(relations) |
| 91 | |
| 92 | x_re = head_re * relation_re - head_im * relation_im - tail_re |
| 93 | x_im = head_re * relation_im + head_im * relation_re - tail_im |
| 94 | x = torch.stack([x_re, x_im], dim=0) |
| 95 | score = margin - x.norm(p=2, dim=0).sum(dim=1) |
| 96 | return score |
| 97 | |
| 98 | @staticmethod |
| 99 | def DistMult(heads, relations, tails): |
| 100 | x = heads * relations * tails |
| 101 | score = x.sum(dim=1) |
| 102 | return score |
no outgoing calls
no test coverage detected