File size: 1,513 Bytes
897fe06 8db8077 897fe06 5a3962c 897fe06 9a9a2f5 897fe06 9a9a2f5 897fe06 8db8077 897fe06 8db8077 9a9a2f5 897fe06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | import torch
import torch.nn as nn
import torch.nn.functional as F
class MLPProjection(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.linear1 = nn.Linear(input_dim, hidden_dim)
self.dropout = nn.Dropout(0.5)
self.linear2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x_output):
# only use first token ([CLS]) of each output
x = x_output[:, 0, :]
x = self.linear1(x)
x = F.silu(x)
x = self.dropout(x)
x = self.linear2(x)
return x
class MLPPrediction(nn.Module):
def __init__(self, input_dim, use_abs_diff=False, use_mult=False):
super().__init__()
self.use_abs_diff = use_abs_diff
self.use_mult = use_mult
real_input_dim = input_dim * (2 + int(use_abs_diff) + int(use_mult))
self.mlp = nn.Sequential(
nn.Linear(real_input_dim, 512),
nn.SiLU(),
nn.Dropout(0.5),
nn.Linear(512, 256),
nn.SiLU(),
nn.Dropout(0.5),
nn.Linear(256, 128),
nn.SiLU(),
nn.Linear(128, 1),
)
def forward(self, x1, x2):
x = torch.cat([x1, x2], dim=1)
if self.use_abs_diff:
x_diff = torch.abs(x1 - x2)
x = torch.cat([x, x_diff], dim=1)
if self.use_mult:
x_mult = x1 * x2
x = torch.cat([x, x_mult], dim=1)
x = self.mlp(x)
return x
|