File size: 2,696 Bytes
5b912c3 | 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 | """
Student model for LoRaVoiceLink: a compact spectrogram U-Net that maps a
Codec2-degraded magnitude spectrogram to an enhanced one.
Kept deliberately small (edge-deployment target: real-time inference on a
Jetson) — unlike the teacher models used only offline during data
preparation, this is the model that actually ships.
"""
import torch
import torch.nn as nn
def conv_block(in_ch, out_ch):
return nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
nn.BatchNorm2d(out_ch),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
nn.BatchNorm2d(out_ch),
nn.LeakyReLU(0.2, inplace=True),
)
class SpectrogramUNet(nn.Module):
"""
Input: (B, 1, F, T) degraded magnitude spectrogram — F and T must be
multiples of 8 (three pooling stages); see dataset.py's
pad_to_multiple().
Output: (B, 1, F, T) enhanced magnitude spectrogram, produced by
predicting a [0, 1] mask and multiplying it against the input
(standard, more stable to train than predicting raw magnitude
directly).
"""
def __init__(self, base_channels: int = 32):
super().__init__()
c = base_channels
self.enc1 = conv_block(1, c)
self.enc2 = conv_block(c, c * 2)
self.enc3 = conv_block(c * 2, c * 4)
self.bottleneck = conv_block(c * 4, c * 8)
self.pool = nn.MaxPool2d(2)
self.up3 = nn.ConvTranspose2d(c * 8, c * 4, kernel_size=2, stride=2)
self.dec3 = conv_block(c * 8, c * 4)
self.up2 = nn.ConvTranspose2d(c * 4, c * 2, kernel_size=2, stride=2)
self.dec2 = conv_block(c * 4, c * 2)
self.up1 = nn.ConvTranspose2d(c * 2, c, kernel_size=2, stride=2)
self.dec1 = conv_block(c * 2, c)
self.out_conv = nn.Conv2d(c, 1, kernel_size=1)
def forward(self, x):
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
b = self.bottleneck(self.pool(e3))
d3 = self.up3(b)
d3 = self.dec3(torch.cat([d3, e3], dim=1))
d2 = self.up2(d3)
d2 = self.dec2(torch.cat([d2, e2], dim=1))
d1 = self.up1(d2)
d1 = self.dec1(torch.cat([d1, e1], dim=1))
mask = torch.sigmoid(self.out_conv(d1))
return mask * x
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
if __name__ == "__main__":
model = SpectrogramUNet()
dummy = torch.randn(2, 1, 256, 128)
out = model(dummy)
print("output shape:", out.shape)
print("parameters:", f"{count_parameters(model):,}")
|