File size: 2,479 Bytes
4dce182 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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())
|