| """ |
| 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):,}") |
|
|
