| """Shared building blocks. Attention is written out by hand so it exports cleanly to ONNX and runs on MPS.""" |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ConvBNAct(nn.Sequential): |
| def __init__(self, cin, cout, k=3, s=1, p=None, groups=1, act=True): |
| if p is None: |
| p = (k // 2) if isinstance(k, int) else tuple(x // 2 for x in k) |
| layers = [nn.Conv2d(cin, cout, k, s, p, groups=groups, bias=False), nn.BatchNorm2d(cout)] |
| if act: |
| layers.append(nn.ReLU(inplace=True)) |
| super().__init__(*layers) |
|
|
|
|
| class SelfAttention(nn.Module): |
| def __init__(self, d, heads, dropout=0.0): |
| super().__init__() |
| assert d % heads == 0 |
| self.h, self.dk = heads, d // heads |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x, bias=None): |
| """x: (B, T, D). bias: additive attention bias broadcastable to (B, H, T, T).""" |
| B, T, D = x.shape |
| qkv = self.qkv(x).reshape(B, T, 3, self.h, self.dk).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.dk)) |
| if bias is not None: |
| att = att + bias |
| att = self.drop(att.softmax(-1)) |
| out = (att @ v).transpose(1, 2).reshape(B, T, D) |
| return self.proj(out) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, d, heads, ff, dropout=0.1): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(d) |
| self.attn = SelfAttention(d, heads, dropout) |
| self.ln2 = nn.LayerNorm(d) |
| self.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(ff, d)) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x, bias=None): |
| x = x + self.drop(self.attn(self.ln1(x), bias)) |
| return x + self.drop(self.ff(self.ln2(x))) |
|
|
|
|
| def sinusoid_table(n, d): |
| pos = torch.arange(n, dtype=torch.float32)[:, None] |
| div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * (-math.log(10000.0) / d)) |
| pe = torch.zeros(n, d) |
| pe[:, 0::2] = torch.sin(pos * div) |
| pe[:, 1::2] = torch.cos(pos * div) |
| return pe |
|
|
|
|
| def padding_bias(pad_mask, dtype): |
| """pad_mask: (B, T) bool, True = padding -> additive bias (B, 1, 1, T).""" |
| return pad_mask[:, None, None, :].to(dtype) * -1e4 |
|
|
|
|
| def count_params(module: nn.Module) -> int: |
| return sum(p.numel() for p in module.parameters()) |
|
|
