File size: 3,433 Bytes
3f792e3 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import torch
import torch.nn as nn
import torch.nn.functional as F
class AttnProjection(nn.Module):
def __init__(self, input_dim, n_heads, output_length):
super().__init__()
self.query = nn.Parameter(torch.randn(output_length, input_dim))
self.attn = nn.MultiheadAttention(
input_dim, n_heads, dropout=0.2, batch_first=True
)
self.norm1 = nn.LayerNorm(input_dim)
self.dropout1 = nn.Dropout(0.2)
self.self_attn = nn.MultiheadAttention(
input_dim, n_heads, dropout=0.2, batch_first=True
)
self.norm2 = nn.LayerNorm(input_dim)
self.dropout2 = nn.Dropout(0.2)
self.cls_mlp = nn.Sequential(
nn.Linear(input_dim, input_dim), nn.SiLU(), nn.Dropout(0.2)
)
self.norm3 = nn.LayerNorm(input_dim)
nn.init.xavier_normal_(self.query)
def forward(self, x):
B = x.shape[0]
query = self.query.unsqueeze(0).repeat(B, 1, 1)
x_cls = x[:, 0, :]
x_other = x[:, 1:, :]
z_other = self.norm1(x_other)
z_attn = self.attn(query, z_other, z_other)[0]
z_other = self.dropout1(z_attn)
z_other = self.norm2(z_other)
z_attn = self.self_attn(z_other, z_other, z_other)[0]
z_other = z_other + self.dropout1(z_attn)
z_cls = x_cls + self.cls_mlp(self.norm3(x_cls))
z = torch.cat([z_cls.unsqueeze(1), z_other], dim=1)
z = z.contiguous().view(B, -1)
return z
class BiAttnPrediction(nn.Module):
def __init__(self, input_dim, n_heads):
super().__init__()
self.input_dim = input_dim
self.attn1 = nn.MultiheadAttention(
input_dim, n_heads, dropout=0.2, batch_first=True
)
self.norm1 = nn.LayerNorm(input_dim)
self.dropout1 = nn.Dropout(0.2)
self.attn2 = nn.MultiheadAttention(
input_dim, n_heads, dropout=0.2, batch_first=True
)
self.norm2 = nn.LayerNorm(input_dim)
self.dropout2 = nn.Dropout(0.2)
self.mlp = nn.Sequential(
nn.Linear(input_dim * 6, 1024),
nn.SiLU(),
nn.Dropout(0.2),
nn.Linear(1024, 512),
nn.SiLU(),
nn.Dropout(0.2),
nn.Linear(512, 256),
nn.SiLU(),
nn.Dropout(0.2),
nn.Linear(256, 1),
)
self.norm3 = nn.LayerNorm(input_dim)
def forward(self, x1, x2):
B = x1.shape[0]
x1 = x1.view(B, -1, self.input_dim) # [B, M x D] -> [B, M, D]
x2 = x2.view(B, -1, self.input_dim) # [B, M x D] -> [B, M, D]
z1_cls = x1[:, 0, :]
z2_cls = x2[:, 0, :]
x1_other = self.norm1(x1[:, 1:, :])
x2_other = self.norm2(x2[:, 1:, :])
z1_attn = self.attn1(x2_other, x1_other, x1_other)[0]
z1_other = x1_other + self.dropout1(z1_attn)
z2_attn = self.attn2(x1_other, x2_other, x2_other)[0]
z2_other = x2_other + self.dropout2(z2_attn)
z1_other = z1_other.mean(dim=1)
z2_other = z2_other.mean(dim=1)
z = torch.cat(
[
z1_cls,
z1_other,
z2_cls,
z2_other,
torch.abs(z1_cls - z2_cls),
torch.abs(z1_other - z2_other),
],
dim=1,
) # [B, D * 4]
z = self.mlp(z)
return z
|