File size: 5,023 Bytes
7df09bd | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """TinyConformer-CTC model, matching the notebook-03 architecture exactly.
Presets:
nano (96, 5): ~1.4M params, ~1.4MB int8
micro (112, 6): ~2.1M params, ~2.1MB int8
tiny (128, 6): ~2.8M params, ~2.8MB int8
small (256, 9): ~16M params, ~16MB int8 ← v2 target
CTC (not RNN-T/seq2seq) so the decode stays trivial. QKV is explicit Linear layers
(not nn.MultiheadAttention) so int8 dynamic quantization covers every big matmul.
Pre-LN on every sublayer (ff1, attn, conv, ff2) for training stability at larger sizes.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class FFN(nn.Module):
def __init__(self, d, mult=4, p=0.1):
super().__init__()
self.net = nn.Sequential(nn.Linear(d, mult*d), nn.SiLU(), nn.Dropout(p),
nn.Linear(mult*d, d), nn.Dropout(p))
def forward(self, x): return self.net(x)
class MHSA(nn.Module):
"""Self-attention with explicit Linear QKV/O so int8 dynamic quant covers it."""
def __init__(self, d, heads=4, p=0.1):
super().__init__()
assert d % heads == 0
self.h, self.dh = heads, d // heads
self.qkv = nn.Linear(d, 3*d)
self.o = nn.Linear(d, d)
self.drop = nn.Dropout(p)
def forward(self, x, pad_mask=None):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, -1)
q = q.view(B, T, self.h, self.dh).transpose(1, 2)
k = k.view(B, T, self.h, self.dh).transpose(1, 2)
v = v.view(B, T, self.h, self.dh).transpose(1, 2)
s = q @ k.transpose(-2, -1) / math.sqrt(self.dh)
if pad_mask is not None:
s = s.masked_fill(pad_mask[:, None, None, :], float("-inf"))
w = self.drop(s.softmax(-1))
return self.o((w @ v).transpose(1, 2).reshape(B, T, D))
class ConvModule(nn.Module):
def __init__(self, d, k=15, p=0.1):
super().__init__()
self.pw1 = nn.Conv1d(d, 2*d, 1)
self.dw = nn.Conv1d(d, d, k, padding=k//2, groups=d)
self.bn = nn.BatchNorm1d(d)
self.pw2 = nn.Conv1d(d, d, 1)
self.drop = nn.Dropout(p)
def forward(self, x): # (B,T,D)
x = x.transpose(1, 2)
x = self.pw1(x)
a, b = x.chunk(2, 1)
x = a * torch.sigmoid(b)
x = self.bn(self.dw(x))
x = self.drop(self.pw2(F.silu(x)))
return x.transpose(1, 2)
class ConformerBlock(nn.Module):
def __init__(self, d=112, heads=4, p=0.1):
super().__init__()
self.ff1_norm = nn.LayerNorm(d)
self.ff1 = FFN(d, p=p)
self.attn_norm = nn.LayerNorm(d)
self.attn = MHSA(d, heads, p)
self.conv_norm = nn.LayerNorm(d)
self.conv = ConvModule(d, p=p)
self.ff2_norm = nn.LayerNorm(d)
self.ff2 = FFN(d, p=p)
self.final = nn.LayerNorm(d)
def forward(self, x, pad_mask=None):
x = x + 0.5 * self.ff1(self.ff1_norm(x))
x = x + self.attn(self.attn_norm(x), pad_mask)
x = x + self.conv(self.conv_norm(x))
x = x + 0.5 * self.ff2(self.ff2_norm(x))
return self.final(x)
class TinyConformerCTC(nn.Module):
N_MELS = 80
def __init__(self, d=112, n_blocks=6, n_classes=29, p=0.1):
super().__init__()
self.c1 = nn.Conv2d(1, d, 3, stride=2, padding=1)
self.c2 = nn.Conv2d(d, d, 3, stride=2, padding=1)
self.proj = nn.Linear(d * (self.N_MELS // 4), d)
self.blocks = nn.ModuleList([ConformerBlock(d, p=p) for _ in range(n_blocks)])
self.head = nn.Linear(d, n_classes)
@staticmethod
def out_len(l): # two stride-2 k3p1 convs
return ((l + 2 - 3)//2 + 1 + 2 - 3)//2 + 1
def forward(self, mel, lens):
x = F.relu(self.c1(mel)); x = F.relu(self.c2(x)) # (B,d,T/4,20)
B, C, T, F_ = x.shape
x = self.proj(x.permute(0, 2, 1, 3).reshape(B, T, C*F_))
# tensor arithmetic keeps out_len ONNX-traceable (avoids baking a constant)
ol = (lens.long().to(x.device) - 1) // 2 + 1
ol = (ol - 1) // 2 + 1
pad = torch.arange(T, device=x.device)[None, :] >= ol[:, None]
for blk in self.blocks:
x = blk(x, pad)
return self.head(x).log_softmax(-1), ol
def count_params(m):
n = sum(p.numel() for p in m.parameters())
print(f"params: {n/1e6:.2f}M | fp32 {n*4/1e6:.1f}MB | int8 ~{n/1e6:.1f}MB (linears)")
return n
def greedy_decode(log_probs, alphabet, blank=0):
"""log_probs: (T, C) -> (text, confidence). Collapses repeats, drops blanks."""
probs = log_probs.exp()
ids = probs.argmax(-1).tolist()
conf = probs.amax(-1)
out, prev, cs = [], None, []
for t, i in enumerate(ids):
if i != blank:
cs.append(float(conf[t]))
if i != prev:
out.append(alphabet[i])
prev = i
text = "".join(out)
return text, (sum(cs)/len(cs) if cs else 0.0)
# confidence < threshold -> "uncertain" (tongue's tie-reporting analog)
UNCERTAIN_TH = 0.35
|