torchaudio 重采样变换缓存内核怎样减少重复计算

音频预处理里最容易被忽略的一段开销，是每个文件都重新设计一次重采样卷积核。固定采样率转换时，应把 torchaudio.transforms.Resample 放到数据循环外创建；它会预计算并缓存内核，后续波形只重复卷积。若每条数据都用 torchaudio.functional.resample()，内核会在调用路径中重新计算，批量处理时尤其浪费。

官方地址：https://docs.pytorch.org/audio/stable/

一、先判断重复计算发生在哪里

functional.resample 适合一次性转换，调用时根据采样率、lowpass_filter_width、rolloff、窗口方式和 beta 等参数生成或准备卷积核；transforms.Resample 则把这一步移到构造阶段，并把内核作为模块状态保存下来。换句话说，真正要复用的是“参数确定后的变换对象”，不是只复用一个函数名。

二、把 Resample 移到数据循环外

下面的写法把初始化放在处理函数外侧。示例中的每条波形共享 16 kHz 到 8 kHz 的同一组参数，循环内只执行变换。

import torch
import torchaudio

# 固定采样率和质量参数，构造阶段预计算重采样内核
resampler = torchaudio.transforms.Resample(
    orig_freq=16_000,
    new_freq=8_000,
    lowpass_filter_width=6,
    rolloff=0.99,
    resampling_method="sinc_interp_hann",
    dtype=torch.float32,  # 用 float32 保存缓存内核，兼顾精度和内存
)

def convert_batch(waveforms):
    # 复用同一个模块；输入通常为 [channel, time]
    return [resampler(waveform) for waveform in waveforms]

不要把 Resample(...) 写进列表推导式，否则每次迭代又回到“重复构造、重复建核”。如果波形已经放在 GPU 上，先把变换模块移动到同一设备；多进程 DataLoader 则应留意每个 worker 都会拥有自己的模块和内核副本。

三、多组采样率用参数元组分层缓存

真实数据集常同时出现 8 kHz、16 kHz 和 48 kHz。此时不能把所有输入硬塞进一个 Resample，而应把“原采样率、目标采样率、质量参数、dtype、设备”组成稳定键。相同键命中缓存，不同键创建各自的内核。

import torch
import torchaudio

_resamplers = {}

def get_resampler(orig_freq, new_freq, device="cpu"):
    # 参数元组决定内核；不要只用 orig_freq 作为缓存键
    key = (orig_freq, new_freq, 6, 0.99, "sinc_interp_hann", "float32", device)
    if key not in _resamplers:
        module = torchaudio.transforms.Resample(
            orig_freq, new_freq,
            lowpass_filter_width=key[2],
            rolloff=key[3],
            resampling_method=key[4],
            dtype=torch.float32,
        ).to(device)
        _resamplers[key] = module
    return _resamplers[key]

def convert(waveform, orig_freq, new_freq):
    # 先按输入实际设备取模块，避免跨设备拷贝或调用失败
    module = get_resampler(orig_freq, new_freq, str(waveform.device))
    return module(waveform)

四、用证据判断缓存是否值得保留

先分别测量“每个样本新建模块”和“循环外复用模块”的平均耗时，再观察内核规模。lowpass_filter_width 越大，滤波更锐利，但内核更大，建核与卷积都更慢；sinc_interp_kaiser 通常比默认 Hann 需要更多计算。采样率比的最大公约数较大时，内核还可能被简化。

缓存命中不等于所有开销消失：长音频卷积、数据搬运和多 worker 的内存占用仍要测量。若只转换一次，直接调用 functional 形式更简单；只有同一组参数被反复使用时，预计算内核的收益才会稳定出现。

常见问题

改变 rolloff 后还能复用原来的 Resample 吗？

不能。rolloff 会改变低通截止位置，也会改变内核；应按新的参数组创建或命中另一项缓存。

为什么复用模块后仍然很慢？

可能是音频很短、建核开销占比本来就低，也可能是输入不断在 CPU 与 GPU 间搬运。把建核、卷积和设备传输分别计时，才能判断缓存是否解决了真正的瓶颈。