Audio Spectrogram Transformer (AST) 入门指南：从零构建音频分类模型

1次阅读

共计 2304 个字符，预计需要花费 6 分钟才能阅读完成。

传统音频分类任务中，卷积神经网络（CNN）一直是主流选择。但当处理长序列音频时，CNN 的局限性逐渐显现：

感受野固定 ：CNN 的卷积核大小决定了其感受野范围，难以捕捉远距离时间维度的依赖关系
平移不变性陷阱 ：音频中的关键特征（如鸟叫声）出现位置可能变化，但 CNN 的平移不变性会模糊这类时序信息
频谱图处理粗糙 ：Mel 频谱图作为二维输入，CNN 往往采用粗暴的二维卷积，忽略了频率轴和时间轴的不同特性

AST 将视觉 Transformer 成功适配到音频领域，核心创新在于：

频谱图分块嵌入 ：将 128×128 的 Mel 频谱图切割为 16×16 的 patch（每 patch 8×8 个点）
Transformer 编码器 ：通过自注意力机制建立全局依赖，比 CNN 更擅长建模长序列
可学习位置编码 ：不同于原始 Transformer 的正弦编码，AST 采用可训练的位置向量

模型类型	参数量	推理速度 (ms)	准确率 (%)
CNN	2.3M	12	78.2
AST	5.7M	18	85.6

import torchaudio

def create_spectrogram(waveform, sample_rate=16000):
    # 计算 Mel 频谱图
    transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=sample_rate,
        n_fft=1024,
        hop_length=160,
        n_mels=128
    )
    spec = transform(waveform)  # [1, 128, 101]

    # 对数压缩并归一化
    spec = torch.log(spec + 1e-9)
    spec = (spec - spec.mean()) / spec.std()
    return spec  # 输出形状 [1, 128, 101]

import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=128, patch_size=8, in_chans=1, embed_dim=768):
        super().__init__()
        self.proj = nn.Conv2d(
            in_chans, embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        self.num_patches = (img_size // patch_size) ** 2

    def forward(self, x):
        x = self.proj(x)  # [B, 768, 16, 16]
        x = x.flatten(2).transpose(1, 2)  # [B, 256, 768]
        return x

class ASTModel(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.patch_embed = PatchEmbed()
        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
        self.pos_embed = nn.Parameter(torch.zeros(1, 256 + 1, 768)  # 256 patches + 1 cls_token
        )
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=768, nhead=12,
                dim_feedforward=3072
            ),
            num_layers=12
        )
        self.head = nn.Linear(768, num_classes)

    def forward(self, x):
        x = self.patch_embed(x)  # [B, 256, 768]
        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)  # [B, 257, 768]
        x = x + self.pos_embed
        x = self.transformer(x)
        x = x[:, 0]  # 取 cls_token 对应的输出
        return self.head(x)