Neural Amp Modeler 中 MRSTFT 损失函数在 MPS 设备上的优化探索

在音频信号处理领域，多分辨率短时傅里叶变换(Multi-Resolution Short-Time Fourier Transform, MRSTFT)损失函数是一种常用的音频质量评估指标。本文将深入探讨在Neural Amp Modeler项目中，如何优化MRSTFT损失函数在苹果Metal Performance Shaders(MPS)设备上的运行性能。

问题背景

MRSTFT损失函数通过比较目标音频和生成音频在不同时间-频率分辨率下的频谱特征差异，能够有效捕捉音频信号的时频特性。然而，在MPS设备上运行时，该损失函数遇到了性能瓶颈。

核心问题在于PyTorch当前版本中aten::angle操作尚未在MPS后端实现。当代码执行到计算复数相位角的步骤时，系统会抛出异常并自动回退到CPU计算，这导致整体性能显著下降。

技术分析

经过深入分析，我们发现MRSTFT损失函数中的相位计算是导致性能问题的关键。具体表现为：

1. 在auraloss库的实现中，默认情况下并不使用相位信息计算损失
2. 但代码仍然会计算相位角度，这触发了MPS不支持的torch.angle操作
3. 即使尝试仅将相位计算部分放在CPU上执行，整体性能仍与完全在CPU上计算相当

解决方案探索

我们考虑了多种优化方案：

方案一：环境变量回退

通过设置PYTORCH_ENABLE_MPS_FALLBACK=1环境变量，可以让不支持的操作自动回退到CPU执行。但测试表明这种方案无法带来性能提升。

方案二：修改auraloss库实现

最新版本的auraloss库已经优化了这一问题，它会在不需要相位信息时跳过相位计算。这从根本上避免了触发不支持的MPS操作。

具体实现上，auraloss通过条件判断来跳过torch.angle调用，当perceptual_weighting参数为0时(默认值)，完全不计算相位相关损失。

实施建议

对于Neural Amp Modeler项目，我们建议：

1. 升级到auraloss库的最新版本(commit 0853e9e或更高)
2. 保持默认参数不变(不使用相位加权)
3. 这样MRSTFT损失函数就能完全在MPS设备上运行，无需CPU回退

未来展望

虽然当前解决方案有效，但从长远看，PyTorch社区需要完善MPS后端的复数运算支持。特别是aten::angle操作的实现将彻底解决这类问题。开发者可以关注PyTorch相关issue的进展，并在新版本发布后重新评估性能优化空间。

通过这次优化实践，我们不仅解决了具体的技术问题，更重要的是建立了一套分析GPU加速问题的通用方法：识别关键瓶颈操作、了解框架支持情况、寻找计算流程优化点。这种思路可以应用于其他类似的性能优化场景。