Audio Spectrogram Transformer 原理解析与实战：如何高效处理音频分类任务

1次阅读

没有评论

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

传统 CNN 在处理音频频谱图时存在明显的瓶颈。频谱图本质上是时间 - 频率二维矩阵，而 CNN 的卷积核设计存在几个关键问题：

固定感受野：3×3 或 5×5 的卷积核难以建模远距离时间依赖（如音乐中的重复段落）
层次化特征丢失：池化操作会逐渐压缩时间维度，导致时序信息模糊化
各向同性处理：标准卷积对时间和频率轴一视同仁，忽略了二者物理意义的差异

这导致在 ESC-50 环境音分类数据集中，CNN 模型的 top- 1 准确率通常卡在 75% 左右难以突破。

Audio Spectrogram Transformer（AST）通过 多头注意力 机制实现了三大改进：

全局建模：每个 patch 都能直接关注全图任意位置
动态权重：根据内容自动调整注意力分布（如重点关注谐波区域）
参数效率：相比 CNN 的逐层局部计算，共享的注意力层更节省参数

实测表明，在参数量相近的情况下：

模型类型	ESC-50 Acc	参数量	FLOPs
ResNet18	76.2%	11M	1.3G
AST	82.7%	9.8M	1.1G

将 $F×T$ 的梅尔频谱图切割为 $N$ 个 $P×P$ 的 patch（通常 P =16）：

$$N = \left\lfloor\frac{F}{P}\right\rfloor \times \left\lfloor\frac{T}{P}\right\rfloor$$

每个 patch 展平后通过线性投影得到 $D$ 维向量：

$$z_p = W_e \cdot \text{Flatten}(x_p) + b_e$$

import torch
import torch.nn as nn

class ASTPatchEmbed(nn.Module):
    def __init__(self, freq=128, time=1024, patch_size=16, dim=768):
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Linear(patch_size*patch_size, dim)  # 投影层

        # 可学习的位置编码（比原版 ViT 多 1 维）self.pos_embed = nn.Parameter(torch.randn(1, (freq//patch_size)*(time//patch_size) + 1, dim)
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))  # 分类令牌

    def forward(self, x):  # x: [B, 1, F, T]
        B, _, F, T = x.shape
        p = self.patch_size

        # 分块处理（关键步骤）x = x.unfold(2, p, p).unfold(3, p, p)  # [B,1,H,W,p,p]
        x = x.permute(0,2,3,1,4,5).flatten(1,2)  # [B,N,1,p,p]
        x = x.flatten(2)  # [B,N,p*p]

        # 投影并添加 CLS 令牌
        x = self.proj(x)  # [B,N,D]
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)

        return x + self.pos_embed  # 添加位置编码

分层解冻：
先冻结所有 Transformer 层，仅训练分类头
按从后往前的顺序逐步解冻 encoder 层

学习率策略：

optimizer = torch.optim.AdamW([{'params': model.head.parameters(), 'lr': 1e-3},
    {'params': model.blocks[-4:].parameters(), 'lr': 5e-5},
    {'params': model.blocks[:-4].parameters(), 'lr': 1e-5}
])

梯度检查点：

from torch.utils.checkpoint import checkpoint

class ASTEncoder(nn.Module):
    def forward(self, x):
        for blk in self.blocks:
            x = checkpoint(blk, x)  # 分段计算梯度
        return x

混合精度训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()