在Accelerate/Swift中对齐Python Log-Mel频谱图数值匹配问题

我正在通过CoreML和Accelerate/AVFoundation将OpenAI的Whisper原生移植到macOS和iOS平台，过程中发现自己实现的Log-Mel频谱图与Whisper的Python代码存在数值差异。已完全按照Whisper的方式拆解Log-Mel频谱图提取步骤，Swift代码已匹配以下环节：

• 原始音频采样
• 转换为Float类型
• 归一化Float转换
• 计算Hamming窗
• 加载预计算的Mel滤波器

但在执行加窗FFT时，数值范围相近但具体数值完全不匹配，并非简单舍入误差，差异集中在：

• STFT过程生成的复数
• 矩阵运算结果

对齐的Python代码 #

window = torch.hann_window(N_FFT).to(audio.device)
stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
magnitudes = stft[:, :-1].abs() ** 2
filters = mel_filters(audio.device, n_mels)
mel_spec = filters @ magnitudes
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
log_spec = (log_spec + 4.0) / 4.0

Swift核心代码片段 #

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)
    
    // insert numFFT/2 samples before and numFFT/2 after so we have a extra numFFT amount to process
    audioFloat.insert(contentsOf: [Float](repeating: 0, count: self.numFFT/2), at: 0)
    audioFloat.append(contentsOf: [Float](repeating: 0, count: self.numFFT/2))
    

    // Split Complex arrays holding the mel spectrogram
    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)

    // we need to create 200 x 3000 matrix of STFTs - note we appear to want to output complex numbers (?)
    for (i) in 0 ..< self.melSampleCount
    {
        // Slice numFFTs every hop count (barf) and make a mel spectrum out of it
        var audioFrame = Array<Float>( audioFloat[ (i * self.hopCount) ..< ( (i * self.hopCount) + self.numFFT) ] )
        
        assert(audioFrame.count == self.numFFT)
        
        // Split Complex arrays holding a single FFT result, which gets appended to the
        var sampleReal:[Float] = [Float](repeating: 0, count: self.numFFT/2)
        var sampleImaginary:[Float] = [Float](repeating: 0, count: self.numFFT/2)
        
        sampleReal.withUnsafeMutableBufferPointer { realPtr in
            sampleImaginary.withUnsafeMutableBufferPointer { imagPtr in
                
                vDSP.multiply(audioFrame,
                              hanningWindow,
                              result: &audioFrame)

                var complexSignal = DSPSplitComplex(realp: realPtr.baseAddress!,
                                                    imagp: imagPtr.baseAddress!)
                       
                audioFrame.withUnsafeBytes { unsafeAudioBytes in
                    vDSP.convert(interleavedComplexVector: [DSPComplex](unsafeAudioBytes.bindMemory(to: DSPComplex.self)),
                                 toSplitComplexVector: &complexSignal)
                }
                
                self.fft.forward(input: complexSignal,
                             output: &complexSignal)
            }
        }

        allSampleReal[i] = sampleReal
        allSampleImaginary[i] = sampleImaginary
    }
    
    // We create flattened  3000 x 200 array of DSPSplitComplex values
    var flattnedReal:[Float] = allSampleReal.flatMap { $0 }
    var flattnedImaginary:[Float] = allSampleImaginary.flatMap { $0 }

    // Take the magnitude squared of the matrix, which results in a Result flat array of 3000 x 200 of real floats
    // Then multiply it with our mel filter bank
    let count = flattnedReal.count
    var magnitudes = [Float](repeating: 0, count: count)
    var melSpectroGram = [Float](repeating: 0, count: 80 * 3000)
    
    flattnedReal.withUnsafeMutableBytes { unsafeReal in
        flattnedImaginary.withUnsafeMutableBytes { unsafeImaginary in
            
            let matrix = [DSPSplitComplex](repeating: DSPSplitComplex(realp: unsafeReal.bindMemory(to: Float.self).baseAddress!,
                                                                      imagp: unsafeImaginary.bindMemory(to: Float.self).baseAddress!),
                                           count: count)
            
            // populate magnitude matrix with magnitudes squared
            vDSP_zvmags(matrix, 1, &magnitudes, 1, vDSP_Length(count))
            
            // transpose magnitudes to get our 200 x 3000 matrix
            vDSP_mtrans(magnitudes, 1, &magnitudes, 1, 3000, 200)
            
            // Matrix A, a MxK sized matrix
            // Matrix B, a KxN sized matrix
            
            // MATRIX A mel filters is 80 rows x 200 columns
            // MATRIX B magnitudes is 3000 x 200
            // MATRIX B is TRANSPOSED to be 200 rows x 3000 columns
            // MATRIX C melSpectroGram is 80 rows x 3000 columns
            
            let M: Int32 = 80 // number of rows in matrix A
            let N: Int32 = 3000 // number of columns in matrix B
            let K: Int32 = 200 // number of columns in matrix A and number of rows in
            
            // matrix multiply magitude squared matrix with our filter bank
            cblas_sgemm(CblasRowMajor,
                        CblasNoTrans,           // Transpose A
                        CblasNoTrans,           //
                        M,                      // M Number of rows in matrices A and C.
                        N,                      // N Number of columns in matrices B and C.
                        K,                      // K Number of columns in matrix A; number of rows in matrix B.
                        1.0,                    // Alpha Scaling factor for the product of matrices A and B.
                        self.melFilterMatrix,   // Matrix A
                        K,                      // LDA The size of the first dimension of matrix A; if you are passing a matrix A[m][n], the value should be m.
                        magnitudes,             // Matrix B
                        N,                      // LDB The size of the first dimension of matrix B; if you are passing a matrix B[m][n], the value should be m.
                        0,                      // Beta Scaling factor for matrix C.
                        &melSpectroGram,        // Matrix C
                        N)                      // LDC The size of the first dimension of matrix C; if you are passing a matrix C[m][n], the value should be m.
            
            var minValue: Float = 1e-10
            var maxValue: Float = 0.0
            var maxIndex: vDSP_Length = 0
            var minIndex: vDSP_Length = 0
            
            let melCount = melSpectroGram.count
            
            // get the current max value
            vDSP_maxvi(melSpectroGram, 1, &maxValue, &maxIndex, vDSP_Length(melCount))
            
            // Clip to a set min value, keeping the current max value
            vDSP_vclip(melSpectroGram, 1, &minValue, &maxValue, &melSpectroGram, 1, vDSP_Length(melCount))
            
            // Take the log base 10
            var melCountInt32:UInt32 = UInt32(melCount)
            vvlog10f(&melSpectroGram, melSpectroGram, &melCountInt32)
            
            // get the new max value
            vDSP_maxvi(melSpectroGram, 1, &maxValue, &maxIndex, vDSP_Length(melCount))
            
            // get the new min value
            vDSP_minvi(melSpectroGram, 1, &minValue, &minIndex, vDSP_Length(melCount))
            
            // emulate: log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
            var newMin = maxValue - 8.0
            
            // Clip to new max and updated min
            vDSP_vclip(melSpectroGram, 1, &newMin, &maxValue, &melSpectroGram, 1, vDSP_Length(melCount))
            
            // Add 4 and Divide by 4
            var four:Float = 4.0
            vDSP_vsadd(melSpectroGram, 1, &four, &melSpectroGram, 1, vDSP_Length(melCount))
            vDSP_vsdiv(melSpectroGram, 1, &four, &melSpectroGram, 1, vDSP_Length(melCount))
            
        }
    }
    
    return melSpectroGram
}

关键疑问点 #

• Python输出的STFT尺寸为201×3001，而根据Whisper参数我的计算结果应为200×3000，原因不明？
• 我是否正确计算了重叠步长？我未将结果合并为单个FFT，且在音频帧前后各填充了200个采样点。
• 我假设PyTorch的STFT输出是3000组长度为200的复数，是否正确？
• 假设vDSP的FFT实现与PyTorch大致匹配，是否成立？
• PyTorch的STFT第0个和最后一个元素的虚部均为0，而我的实现并非如此，原因是什么？

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