File size: 932 Bytes
69aac37 | 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 | import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
class PodosTransformer(nn.Module,PyTorchModelHubMixin):
def __init__(self, input_dim, model_dim, num_classes, num_heads=4, num_layers=2, dropout=0.1,temperature=1):
super(PodosTransformer, self).__init__()
self.temperature = temperature
self.projection = nn.Linear(input_dim, model_dim)
encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads, dropout=dropout)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.fc = nn.Linear(model_dim, num_classes)
def forward(self, x):
x = self.projection(x)
x = x.unsqueeze(1)
x = self.transformer_encoder(x)
x = x.mean(dim=1)
x = self.fc(x)
if self.temperature != 1.0:
x = x / self.temperature
return x |