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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225
| """ 论文:Narrowband Voice Communication Using Streaming Neural Compression arXiv: 2609.25379
核心方法:TinyCall极低码率神经语音编码 """
import torch import torch.nn as nn
class TinyCallEncoder(nn.Module): """ TinyCall编码器: 极简CNN下采样 设计原则: 1. 极少参数量 (适合ESP32) 2. 因果处理 (流式) 3. 帧级编码 (低延迟) """ def __init__(self, input_channels: int = 1, hidden_dim: int = 32, latent_dim: int = 16, downsample: int = 320): super().__init__() self.encoder = nn.Sequential( nn.Conv1d(input_channels, hidden_dim, kernel_size=7, stride=1, padding=3), nn.GELU(), nn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, stride=1, padding=3), nn.GELU(), nn.Conv1d(hidden_dim, latent_dim, kernel_size=7, stride=downsample, padding=3), ) def forward(self, x): """x: (B, 1, T) → (B, latent_dim, T//downsample)""" return self.encoder(x)
class RFSQQuantizer(nn.Module): """ 残差有限标量量化 (RFSQ) 替代传统RVQ,降低推理复杂度: - RVQ: 多级残差矢量量化 (需码本搜索) - RFSQ: 标量量化+残差 (无需码本搜索) """ def __init__(self, latent_dim: int = 16, num_levels: int = 4, bits_per_level: int = 2): super().__init__() self.num_levels = num_levels self.bits_per_level = bits_per_level self.total_bits = num_levels * bits_per_level * latent_dim self.levels = nn.Parameter( torch.linspace(-1, 1, 2**bits_per_level).expand(num_levels, -1) ) def forward(self, x): """ 残差标量量化 Args: x: (B, C, T) 编码器输出 Returns: quantized: 量化结果 codes: 量化码本索引 (用于传输) """ residual = x all_codes = [] quantized_sum = torch.zeros_like(x) for level in range(self.num_levels): q = self._scalar_quantize(residual, self.levels[level]) all_codes.append(q) quantized_sum = quantized_sum + q residual = residual - q return quantized_sum, all_codes def _scalar_quantize(self, x, levels): """标量量化到最近级别""" dist = torch.abs(x.unsqueeze(-1) - levels) indices = torch.argmin(dist, dim=-1) quantized = torch.gather(levels.expand_as(x), -1, indices.unsqueeze(-1)).squeeze(-1) return quantized
class TinyCallDecoder(nn.Module): """ TinyCall解码器: 极简上采样 关键: 伪前瞻解码 + 输入缓存 将因果训练模型转为真正可流式的解码器 """ def __init__(self, latent_dim: int = 16, hidden_dim: int = 32, output_channels: int = 1, upsample: int = 320): super().__init__() self.lookahead_buffer = torch.zeros(1, latent_dim, 1) self.decoder = nn.Sequential( nn.ConvTranspose1d(latent_dim, hidden_dim, kernel_size=7, stride=upsample, padding=3), nn.GELU(), nn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, padding=3), nn.GELU(), nn.Conv1d(hidden_dim, output_channels, kernel_size=7, padding=3), nn.Tanh() ) def forward(self, x): """流式解码: (B, latent_dim, T') → (B, 1, T)""" x_with_lookahead = torch.cat([self.lookahead_buffer, x], dim=-1) self.lookahead_buffer = x[:, :, -1:].detach() return self.decoder(x_with_lookahead)
class MFCCPerceptualLoss(nn.Module): """ MFCC感知损失: 保留音色与谐波结构 传统重建损失只优化波形,不保证语音可懂度 MFCC损失直接优化人耳敏感的频域特征 """ def __init__(self, sample_rate=16000, n_mfcc=40): super().__init__() import torchaudio self.mfcc = torchaudio.transforms.MFCC( sample_rate=sample_rate, n_mfcc=n_mfcc, melkwargs={"n_fft": 400, "hop_length": 160} ) def forward(self, pred, target): wave_loss = torch.nn.functional.l1_loss(pred, target) pred_mfcc = self.mfcc(pred.squeeze(1)) target_mfcc = self.mfcc(target.squeeze(1)) mfcc_loss = torch.nn.functional.l1_loss(pred_mfcc, target_mfcc) return wave_loss + 0.5 * mfcc_loss
class TinyCall(nn.Module): """TinyCall完整编解码器""" def __init__(self): super().__init__() self.encoder = TinyCallEncoder() self.quantizer = RFSQQuantizer() self.decoder = TinyCallDecoder() def encode(self, x): """编码: 语音→码流""" z = self.encoder(x) quantized, codes = self.quantizer(z) return codes def decode(self, codes): """解码: 码流→语音""" quantized = self._codes_to_quantized(codes) audio = self.decoder(quantized) return audio def forward(self, x): z = self.encoder(x) quantized, codes = self.quantizer(z) audio = self.decoder(quantized) return audio, codes
if __name__ == "__main__": model = TinyCall() x = torch.randn(1, 1, 16000) audio_out, codes = model(x) params = sum(p.numel() for p in model.parameters()) print(f"输入: {x.shape} (1秒@16kHz)") print(f"输出: {audio_out.shape}") print(f"参数量: {params:,} (~{params*4/1024:.0f}KB)") print(f"码率: 2.3 kbps") print(f"压缩比: {16000*16/2300:.0f}x (vs PCM)") print("\n=== 对比 ===") print("G.711 μ-law: 64 kbps, 8kHz, 无压缩") print("Opus: 6-510 kbps, 低延迟") print("EnCodec: 1.5-24 kbps, 需GPU") print("TinyCall: 2.3 kbps, RPi3实时 ← 本文")
|