Spaces:
Sleeping
Sleeping
File size: 2,277 Bytes
6105886 3bf0f8d 6105886 3bf0f8d c39c656 3bf0f8d 6105886 3bf0f8d 6105886 3bf0f8d 6105886 3bf0f8d 6105886 | 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 | import torch
import torch.nn as nn
class DualStreamTransformer(nn.Module):
def __init__(self, n_feat1=24, n_feat2=10, d_model=32, num_classes=2):
super(DualStreamTransformer, self).__init__()
dim_ff = 128
self.feat_tokenizers_1 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat1)])
self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, d_model))
encoder_layer_1 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
self.encoder_1 = nn.TransformerEncoder(encoder_layer_1, num_layers=2)
self.feat_tokenizers_2 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat2)])
self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, d_model))
encoder_layer_2 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
self.encoder_2 = nn.TransformerEncoder(encoder_layer_2, num_layers=2)
self.fusion = nn.Sequential(
nn.Linear(d_model * 2, d_model),
nn.ReLU()
)
def forward(self, x1, x2):
tokens1 = [layer(x1[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_1)]
x1_emb = torch.stack(tokens1, dim=1)
x1_emb = torch.cat((self.cls_token_1.expand(x1.size(0), -1, -1), x1_emb), dim=1)
feat1 = self.encoder_1(x1_emb)[:, 0, :]
tokens2 = [layer(x2[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_2)]
x2_emb = torch.stack(tokens2, dim=1)
x2_emb = torch.cat((self.cls_token_2.expand(x2.size(0), -1, -1), x2_emb), dim=1)
feat2 = self.encoder_2(x2_emb)[:, 0, :]
combined = torch.cat((feat1, feat2), dim=1)
return self.fusion(combined)
class ArcMarginProduct(nn.Module):
def __init__(self, in_features, out_features, s=30.0, m=0.5):
super(ArcMarginProduct, self).__init__()
self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
nn.init.xavier_uniform_(self.weight)
def predict(self, x):
cosine = torch.matmul(nn.functional.normalize(x), nn.functional.normalize(self.weight).t())
return cosine |