File size: 1,990 Bytes
23ee868 |
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 |
import torch
import torch.nn as nn
from .PositionalEncoding import PositionalEncoding
class AntiCheatPT_256(nn.Module):
def __init__(
self,
feature_dim=44, # nr of features per tick
seq_len=256, # nr of ticks
nhead=1, # nr of attention heads
num_layers=4, # nr of transformer encoder layers
dim_feedforward=176, # hidden size of feedforward network (MLP)
dropout=0.1 # dropout rate
):
super(AntiCheatPT_256, self).__init__()
self.positional_encoding = PositionalEncoding(d_model=feature_dim, max_len=seq_len + 1) # +1 for CLS token
encoder_layer = nn.TransformerEncoderLayer(
d_model=feature_dim,
nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout,
batch_first=True
) # input shape is (batch, seq_len, d_model)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.cls_token = nn.Parameter(torch.zeros(1, 1, feature_dim)) # add classification token
self.fc_out = nn.Sequential(
nn.Linear(feature_dim, 128),
nn.ReLU(),
nn.Linear(128, 1)
)
def forward(self, x):
# x shape: (batch_size, seq_len, feature_dim) = (batch_size, 256, d_model)
B = x.size(0)
# add classification token
cls_tokens = self.cls_token.expand(B, -1, -1).to(x.device) # (batch_size, 1, d_model)
x = torch.cat((cls_tokens, x), dim=1) # -> (batch_size, 257, d_model)
x = self.positional_encoding(x) # add positional encoding
x = self.transformer_encoder(x) # -> (batch_size, 257, d_model)
cls_output = x[:, 0] # get output for classification token
out = self.fc_out(cls_output) # -> (batch_size, 1)
return out.squeeze(1) # -> (batch_size,) |