torchlibrosaとtorchaudioの実行速度比較したら20倍も差があった

ふと、torchlibrosaとtorchaudioの実行速度の違いが気になったので検証しました。

こちらのサイトを参考にさせて頂きました。

https://tam5917.hatenablog.com/entry/2021/03/23/220351

比較に用いたコードはこちら

import torchaudio
import torchaudio.transforms as T
import torchlibrosa.stft as tl
import time
N_FFT = 2048
WIN_LENGTH = None
HOP_LENGTH = 512
N_MELS = 256
sample_rate = 32000
# GPUを使用したときの実行速度を比較
# torchaudio (GPU)
wave_list = torch.rand(120, 64000).cuda()
torch.cuda.synchronize()
start1 = time.time()
mel_spectrogram = T.MelSpectrogram(
    sample_rate=sample_rate,
    n_fft=N_FFT,
    win_length=WIN_LENGTH,
    hop_length=HOP_LENGTH,
    power=2.0,
    n_mels=N_MELS,
).cuda()
melspe = mel_spectrogram(wave_list)
torch.cuda.synchronize()
end1 = time.time()
elapsed_time1 = end1 - start1
print(f"elapsed_time (torchaudio;   GPU): {elapsed_time1:.6f}" + f"[sec], melspe: {melspe.shape}")


# torchlibrosa (GPU)
# TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()
torch.cuda.synchronize()
start2 = time.time()
feature_extractor = torch.nn.Sequential(
    tl.Spectrogram(
        n_fft=N_FFT,
        win_length=WIN_LENGTH,
        hop_length=HOP_LENGTH,
        power=2.0,
    ),
    tl.LogmelFilterBank(
        n_fft=N_FFT,
        sr=sample_rate,
        n_mels=N_MELS,
        is_log=False,  # Default is true
    ),
).cuda()
mel_spe2 = feature_extractor(wave_list)
torch.cuda.synchronize()
end2 = time.time()
elapsed_time2 = end2 - start2
print(f"elapsed_time (torchlibrosa; GPU): {elapsed_time2:.6f}" + f"[sec], melspe:{mel_spe2.shape}")

elapsed_time (torchaudio;   GPU): 0.078266[sec], melspe: torch.Size([132, 256, 1251])
elapsed_time (torchlibrosa; GPU): 1.639459[sec], melspe:torch.Size([132, 1, 1251, 256])

結果からなんと、torchaudioの方がtorchlibrosaよりも、20倍も速いことが分かりました。

音処理をする場合、mixupやtime stretchなど波形に対してaugmentationしてからスペクトログラムに変換したい場合が多いと思います。

on the flyな処理をしたいときにdataloader周りの処理がボトルネックになって実行時間が長いのが気になっていました。

torchaudioを使うことで高速化が期待できそうです。

今までPANNsの実装はtorchlibrosaでやっていましたが、torchaudioに変えたら結構高速化するかもしれません。