File size: 1,356 Bytes
8015fc7 | 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 | from __future__ import annotations
import torch
from schema import ACTIONS, VOCABULARY_SIZE
from torch import nn
class KernelMindPolicy(nn.Module):
def __init__(self, dimensions: int = 32) -> None:
super().__init__()
self.token_embedding = nn.Embedding(VOCABULARY_SIZE, dimensions)
self.class_token = nn.Parameter(torch.zeros(1, 1, dimensions))
self.position_embedding = nn.Parameter(torch.randn(1, 8, dimensions) * 0.02)
layer = nn.TransformerEncoderLayer(
d_model=dimensions,
nhead=4,
dim_feedforward=64,
dropout=0.0,
batch_first=True,
activation="gelu",
norm_first=True,
)
self.encoder = nn.TransformerEncoder(layer, num_layers=2)
self.action_heads = nn.Linear(dimensions, 3 * len(ACTIONS))
def forward(self, tokens: torch.Tensor) -> torch.Tensor:
embedded = self.token_embedding(tokens)
class_token = self.class_token.expand(len(tokens), -1, -1)
sequence = torch.cat([class_token, embedded], dim=1)
hidden = self.encoder(sequence + self.position_embedding)
return self.action_heads(hidden[:, 0]).reshape(len(tokens), 3, len(ACTIONS))
def parameter_count(model: nn.Module) -> int:
return sum(parameter.numel() for parameter in model.parameters())
|