File size: 1,501 Bytes
08e7e4c be943d7 08e7e4c be943d7 08e7e4c be943d7 08e7e4c be943d7 08e7e4c be943d7 08e7e4c | 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 | import torch, torch.nn as nn, math
class SinusoidalPE(nn.Module):
def __init__(self, d_model, max_len=512):
super().__init__()
pe = torch.zeros(max_len, d_model); pos = torch.arange(0, max_len).unsqueeze(1).float()
div = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(pos * div); pe[:, 1::2] = torch.cos(pos * div)
self.register_buffer("pe", pe.unsqueeze(0))
def forward(self, x): return x + self.pe[:, :x.size(1)]
class TransSLR(nn.Module):
def __init__(self, feat_dim=225, d_model=128, nhead=4, num_layers=2,
ffn_dim=512, dropout=0.6, num_classes=30):
super().__init__()
self.proj = nn.Linear(feat_dim, d_model); self.pos_enc = SinusoidalPE(d_model)
enc_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
dim_feedforward=ffn_dim, dropout=dropout, batch_first=True)
self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
self.gap = nn.AdaptiveAvgPool1d(1); self.drop = nn.Dropout(dropout)
self.head = nn.Linear(d_model, num_classes)
def forward(self, x):
x = self.proj(x); x = self.pos_enc(x); x = self.encoder(x)
x = self.gap(x.transpose(1,2)).squeeze(-1)
return self.head(self.drop(x))
def get_embeddings(self, x):
x = self.proj(x); x = self.pos_enc(x); x = self.encoder(x)
return self.gap(x.transpose(1,2)).squeeze(-1)
|