| 487 | |
| 488 | |
| 489 | class BLMLP(nn.Module): |
| 490 | def __init__(self): |
| 491 | super(BLMLP, self).__init__() |
| 492 | self.W = nn.Parameter(nn.init.xavier_uniform_(torch.empty(args.student_embed_size, args.student_embed_size))) |
| 493 | self.act = nn.LeakyReLU(negative_slope=0.5) |
| 494 | |
| 495 | def forward(self, embeds): |
| 496 | pass |
| 497 | |
| 498 | def featureExtract(self, embeds): |
| 499 | return self.act(embeds @ self.W) + embeds |
| 500 | |
| 501 | def pairPred(self, embeds1, embeds2): |
| 502 | return (self.featureExtract(embeds1) * self.featureExtract(embeds2)).sum(dim=-1) |
| 503 | |
| 504 | def crossPred(self, embeds1, embeds2): |
| 505 | return self.featureExtract(embeds1) @ self.featureExtract(embeds2).T |
| 506 | |
| 507 | |
| 508 | |