TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

pocketfftを使ってみる

先日、Whisperで音声認識する処理を実装した際に、FFTの処理でフレームサイズが2のべき乗でない場合にどうやって処理するか戸惑った。

Whisperではフレームサイズが400となっており、NAudioのFFTでは、フレームサイズが2のべき乗という制約があり使用できなかった。

FFTアルゴリズムは、Nが2のべき乗でなくても素因数分解できれば高速に計算できるが、2のべき乗のみを実装しているライブラリが多い。
C#で2のべき乗以外にも対応しているライブラリを探したところ、Math.NET Numericsが対応していたので使用した。

WhisperをC++でも使いたいと考えており、C++で、2のべき乗以外に対応しているライブラリだと何があるのか調べてみた。
なお、私が開発しているスマホアプリでは、FFTの実装に、大浦版FFTを使用しているが、これも2のべき乗である必要がある。

C++で2のべき乗以外に対応しているライブラリ

FFTWは、GPLライセンスのため、あまり使いたくない。
FFTPACKはFortranの実装のラッパーになっている。
pocketfftは、FFTPACKを改良してC++でヘッダーのみで実装されており、Numpyでも採用されている。ライセンスもBSDで使いやすい。

この中では、pocketfftが一番良さそうである。

pocketfftの使い方

説明がほとんどないため、pocketfft_demo.ccのコードを参考にした。

入力が実数で、出力が複素数の場合のコード例

サンプリング周波数16kHzで、800Hz、5000Hzの正弦波と20000Hzの余弦波を合成した波形にFFTする処理である。
フレームサイズは400とする。

#include <complex>
#include <cmath>
#include <iostream>
#include "pocketfft_hdronly.h"

using namespace pocketfft;

constexpr double PI = 3.14159;
constexpr double SAMPLE_RATE = 160000;
constexpr int N_FFT = 400;

int main()
{
	double data[N_FFT];
	std::complex<double> res[N_FFT / 2 + 1];

	for (int i = 0; i < N_FFT; i++) {
		data[i] = sin(2.0 * PI * 800.0 * (double)i / SAMPLE_RATE) + sin(2.0 * PI * 5000.0 * (double)i / SAMPLE_RATE) + cos(2.0 * PI * 20000.0 * (double)i / SAMPLE_RATE);
	}

	shape_t shape{ N_FFT };
	stride_t stride_in(shape.size());
	stride_t stride_out(shape.size());
	size_t tmp_in = sizeof(double);
	size_t tmp_out = sizeof(std::complex<double>);
	for (int i = shape.size() - 1; i >= 0; --i)
	{
		stride_in[i] = tmp_in;
		tmp_in *= shape[i];
		stride_out[i] = tmp_out;
		tmp_out *= shape[i];
	}
	shape_t axes;
	for (size_t i = 0; i < shape.size(); ++i)
		axes.push_back(i);

	r2c(shape, stride_in, stride_out, axes, FORWARD, data, res, 1.);

	std::cout << "index\treal\timag\tabs\n";
	for (int i = 0; i < N_FFT / 2 + 1; i++)
		std::cout << i << "\t" << res[i].real() << "\t" << res[i].imag() << "\t" << std::abs(res[i]) << "\n";

	return 0;
}

stride_in、stride_outの設定はfor文になっているが、入力が1次元の場合はfor文にする必要はないが、サンプルコードのまま記載した。

出力の絶対値をグラフにすると、以下のようになる。

3つのピークが確認できる。

torch.stftと比較

Whisperで使用されているPyTorchのtorch.stftと結果を比較した。

import torch
import math

SAMPLE_RATE = 160000

d = torch.zeros(400)
for i in range(400):
    d[i] = math.sin(2*math.pi*800*i/SAMPLE_RATE) + math.sin(2*math.pi*5000*i/SAMPLE_RATE) + math.cos(2*math.pi*20000*i/SAMPLE_RATE)

stft = torch.stft(d, 400, center=False, return_complex=True)
pd.DataFrame(stft.abs().squeeze().numpy()).plot()


同じ結果が得られた。

まとめ

フレームサイズが2のべき乗以外に対応しているC++FFTライブラリを調べた。
pocketfftが、Numpyでも採用されており、ヘッダーのみで、ライブラリもBSDで扱いことがわかった。
pocketfftを使用して、実数入力複素数出力でFFTするコードを実装して試してみた。