Swift用Accelerate复现Torch STFT遇结果偏差，求技术解惑

背景 #

我正在Swift环境下使用Accelerate/vDSP库重实现PyTorch的STFT代码，目标是生成Log Mel频谱图作为OpenAI Whisper模型CoreML移植版本的输入。

结果差异 #

PyTorch原生STFT/Mel代码生成的频谱图（注：原始float32数据导入Photoshop导致被截断）：

我生成的频谱图：

两者在数值和频率分量偏移上存在明显差异。

参考依据 #

PyTorch的STFT官方文档：torch.stft

STFT公式：

$$X[\omega,m]= \sum_{k=0}^{win_length-1} window[k] \cdot input[m×hop_length+k] \cdot exp\left(−j \cdot \frac{2π⋅ωk}{win_length}\right)$$

困惑点 #

我认为自己已正确处理window[k] * input[m×hop_length+k]部分，但对以下内容存在疑问：

• 公式中的指数计算逻辑
• -j的含义（复数虚部单位）
• 如何在vDSP中对应实现该指数运算
• 如何通过求和得到所需的200个元素结果

我的实现代码 #

func processData(audio: [Int16]) -> [Float]
{
    assert(self.sampleCount == audio.count)
        
    var audioFloat:[Float] = [Float](repeating: 0, count: audio.count)
            
    vDSP.convertElements(of: audio, to: &audioFloat)
    
    vDSP.divide(audioFloat, 32768.0, result: &audioFloat)
    
    // 到这一步，Python和Swift的数值完全一致
         
    // 在前后各插入numFFT/2个0样本，凑出额外的numFFT长度用于处理
    // TODO: 这一步是必须的吗？
    audioFloat.insert(contentsOf: [Float](repeating: 0, count: self.numFFT/2), at: 0)
    audioFloat.append(contentsOf: [Float](repeating: 0, count: self.numFFT/2))

    // 存储FFT结果的拆分复数数组
    var allSampleReal = [[Float]](repeating: [Float](repeating: 0, count: self.numFFT/2), count: self.melSampleCount)
    var allSampleImaginary = [[Float]](repeating: [Float](repeating: 0, count: self.numFFT/2), count: self.melSampleCount)

    
    // 步骤2 - 需要创建200×3000的STFT矩阵 - 注意输出应该是复数？
    for (m) in 0 ..< self.melSampleCount
    {
        // 每隔hopCount取一段numFFT长度的样本，用于生成梅尔频谱
        // audioFrame最终存放拆分后的复数
        var audioFrame = Array<Float>( audioFloat[ (m * self.hopCount) ..< ( (m * self.hopCount) + self.numFFT) ] )
        
        // audioFrame原始样本的副本
        let audioFrameOriginal = audioFrame
        
        assert(audioFrame.count == self.numFFT)
        
        // 存储单帧音频FFT结果的拆分复数数组，之后会追加到allSample系列数组中
        var sampleReal:[Float] = [Float](repeating: 0, count: self.numFFT/2)
        var sampleImaginary:[Float] = [Float](repeating: 0, count: self.numFFT/2)

        sampleReal.withUnsafeMutableBytes { unsafeReal in
            sampleImaginary.withUnsafeMutableBytes { unsafeImaginary in
                
                vDSP.multiply(audioFrame,
                              hanningWindow,
                              result: &audioFrame)

                var complexSignal = DSPSplitComplex(realp: unsafeReal.bindMemory(to: Float.self).baseAddress!,
                                                    imagp: unsafeImaginary.bindMemory(to: Float.self).baseAddress!)
                       
                audioFrame.withUnsafeBytes { unsafeAudioBytes in
                    vDSP.convert(interleavedComplexVector: [DSPComplex](unsafeAudioBytes.bindMemory(to: DSPComplex.self)),
                                 toSplitComplexVector: &complexSignal)
                }
                
                // 步骤3 - 执行FFT
                self.fft.forward(input: complexSignal, output: &complexSignal)                    
            }
        }

        // 需要对齐PyTorch的torch.stft实现
        // 到这里我不知道该怎么继续了？
        
//        let twoπ = Float.pi * 2
//        let freqstep:Float = Float(16000 / (self.numFFT/2))
//
//        var w:Float = 0.0
//        for (k) in 0 ..< self.numFFT/2
//        {
//            let j:Float = sampleImaginary[k]
//            let sample = audioFrame[k]
//
//            let exponent = -j * ( (twoπ * freqstep * Float(k) ) / Float((self.numFFT/2)))
//
//            w += powf(sample, exponent)
//        }
        
        
        allSampleReal[m] = sampleReal
        allSampleImaginary[m] = sampleImaginary
    }
    
    // 现在allSample系列拆分复数数组中存储了3000个200维的实部和虚部FFT结果
    
    // 将其展平为3000×200的DSPSplitComplex值数组
    var flattnedReal:[Float] = allSampleReal.flatMap { $0 }
    var flattnedImaginary:[Float] = allSampleImaginary.flatMap { $0 }

}

内容的提问来源于stack exchange，提问作者vade