| 69 | return self.relu(torch.mm(x_direction, self.weight.T) + self.bias.unsqueeze(0)) |
| 70 | |
| 71 | def train(self, x_pos, x_neg): |
| 72 | for i in range(self.num_epochs): |
| 73 | g_pos = self.forward(x_pos).pow(2).mean(1) |
| 74 | g_neg = self.forward(x_neg).pow(2).mean(1) |
| 75 | loss = torch.log1p( |
| 76 | torch.exp( |
| 77 | torch.cat([-g_pos + self.threshold, g_neg - self.threshold]) |
| 78 | ) |
| 79 | ).mean() |
| 80 | self.opt.zero_grad() |
| 81 | loss.backward() |
| 82 | self.opt.step() |
| 83 | if i % args.log_interval == 0: |
| 84 | print("Loss: ", loss.item()) |
| 85 | return self.forward(x_pos).detach(), self.forward(x_neg).detach() |
| 86 | |
| 87 | |
| 88 | if __name__ == "__main__": |