torchaudioとtorchlibrosaの実行速度に違いはあるのか？

はじめに

PyTorchには音声系データを処理するのに便利なtorchaudioというライブラリが存在する。
pytorch.org

一方、音声系データの処理に便利なlibrosaというパッケージが存在する。
librosa.org

さらにtorchlibrosaという、librosa内部の行列計算まわりをPyTorchで置き換えたパッケージが存在する。
github.com

ここで一つ疑問：
「で、結局どれ使えばいいの？（どれが早いの？）」

この疑問が気になったので、CPU実行とGPU実行における実行時間を比較検証するための簡易的なスクリプトを書いて調べてみたということである。

環境

OS、ハードウェア、ドライバ
- Ubuntu 18.04
- GeForce RTX 2070
- CUDA 11.2
- CUDA Driver 460.32.03

ソフトウェア
- Python 3.6.9
- librosa 0.8.0
- torch 1.7.1
- torchaudio 0.7.2
- torchlibrosa 0.0.9

スクリプト

メルスペクトログラムを計算するスクリプトである。

3月24日追記：GPUの計測方法に誤りがあったため、コメントに従ってtorch.cuda.synchronize()を用いるように修正した。

スクリプトを表示する

"""
PyTorch script for verification of execution time.
"""

# Commentary:
# torch 1.7.1
# torchaudio 0.7.2
# torchlibrosa 0.0.9

# GeForce RTX 2070
# CUDA 11.2
# CUDA Driver 460.32.03

import time

import librosa
import librosa.feature
import torch
import torchaudio
import torchaudio.transforms as T
import torchlibrosa as tl

N_FFT = 1024
WIN_LENGTH = None
HOP_LENGTH = 512
N_MELS = 128

if __name__ == "__main__":

    waveform, sample_rate = librosa.load(librosa.ex("nutcracker"))
    # librosa
    start = time.time()
    mel_spectrogram = librosa.feature.melspectrogram(
        y=waveform,
        sr=sample_rate,
        n_fft=N_FFT,
        hop_length=HOP_LENGTH,
        n_mels=N_MELS,
        power=2.0,
    )
    elapsed_time = time.time() - start
    print("elapsed_time (librosa): {0:.6f}".format(elapsed_time) + "[sec]")

    waveform, sample_rate = torchaudio.load(librosa.ex("nutcracker"))
    waveform_cuda = waveform.cuda()

    # torchaudio (CPU)
    start = 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,
    )
    melspec = mel_spectrogram(waveform)
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; CPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchaudio (GPU)
    torch.cuda.synchronize()
    start = 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,
    )
    mel_spectrogram = mel_spectrogram.cuda()
    melspec = mel_spectrogram(waveform_cuda)
    torch.cuda.synchronize()
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; GPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchlibrosa (CPU)
    # TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()
    start = 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
        ),
    )
    logmel = feature_extractor(waveform)
    elapsed_time = time.time() - start
    print("elapsed_time (torchlibrosa; CPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchlibrosa (GPU)
    # TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()
    torch.cuda.synchronize()
    start = 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
        ),
    )
    feature_extractor = feature_extractor.cuda()
    logmel = feature_extractor(waveform_cuda)
    torch.cuda.synchronize()
    elapsed_time = time.time() - start
    print("elapsed_time (torchlibrosa; GPU): {0:.6f}".format(elapsed_time) + "[sec]")

5月8日追記：koukyo1213様のコメントに従い、torch.nn.Sequentialやtorchaudio.transforms.MelSpectrogramのインスタンス化の部分を計測対象から外し、計算処理のみにフォーカスしたスクリプトは以下である。

スクリプトを表示する

"""
PyTorch script for verification of execution time.
"""

# Commentary:
# torch 1.8.1
# torchaudio 0.8.1
# torchlibrosa 0.0.9

# GeForce RTX 2070
# CUDA 11.2
# CUDA Driver 460.32.03

import time

import librosa
import librosa.feature
import torch
import torchaudio
import torchaudio.transforms as T
import torchlibrosa as tl

N_FFT = 1024
WIN_LENGTH = None
HOP_LENGTH = 512
N_MELS = 128

if __name__ == "__main__":

    torchaudio.set_audio_backend("sox_io")

    waveform, sample_rate = librosa.load(librosa.ex("nutcracker"))
    # librosa
    start = time.time()
    mel_spectrogram = librosa.feature.melspectrogram(
        y=waveform,
        sr=sample_rate,
        n_fft=N_FFT,
        hop_length=HOP_LENGTH,
        n_mels=N_MELS,
        power=2.0,
    )
    elapsed_time = time.time() - start
    print("elapsed_time (librosa): {0:.6f}".format(elapsed_time) + "[sec]")

    waveform, sample_rate = torchaudio.load(librosa.ex("nutcracker"))
    waveform_cuda = waveform.cuda()

    # torchaudio (cpu)
    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,
    )
    start = time.time()
    melspec = mel_spectrogram(waveform)
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; CPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchaudio (GPU)
    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,
    )
    torch.cuda.synchronize()
    start = time.time()
    mel_spectrogram = mel_spectrogram.cuda()
    melspec = mel_spectrogram(waveform_cuda)
    torch.cuda.synchronize()
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; GPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchlibrosa (CPU)
    # TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()
    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
        ),
    )
    start = time.time()
    logmel = feature_extractor(waveform)
    elapsed_time = time.time() - start
    print("elapsed_time (torchlibrosa; CPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchlibrosa (GPU)
    # TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()
    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
        ),
    )
    torch.cuda.synchronize()
    start = time.time()
    feature_extractor = feature_extractor.cuda()
    logmel = feature_extractor(waveform_cuda)
    torch.cuda.synchronize()
    elapsed_time = time.time() - start
    print("elapsed_time (torchlibrosa; GPU): {0:.6f}".format(elapsed_time) + "[sec]")

実行結果

その1（3月24日版）

GPUを使ったtorchaudioが一番早い結果となった。

elapsed_time (librosa): 0.042655[sec]
elapsed_time (torchaudio; CPU): 0.021293[sec]
elapsed_time (torchaudio; GPU): 0.002648[sec]
elapsed_time (torchlibrosa; CPU): 0.237388[sec]
elapsed_time (torchlibrosa; GPU): 0.209118[sec]

その2（5月8日版）

koukyo1213様のコメントに従って、torch.nn.Sequentialやtorchaudio.transforms.MelSpectrogramのインスタンス化部分を計測対象から外した場合の結果を以下に示す。

elapsed_time (librosa): 0.040965[sec]
elapsed_time (torchaudio; CPU): 0.014024[sec]
elapsed_time (torchaudio; GPU): 0.001927[sec]
elapsed_time (torchlibrosa; CPU): 0.027673[sec]
elapsed_time (torchlibrosa; GPU): 0.006538[sec]

今回もGPUを使ったtorchaudioが一番早い結果となった。

考察とまとめ

最初のスクリプトでtorchlibrosaの実行時間がより多くかかったのは、koukyo1213様ご指摘の通り、インスタンス化のオーバーヘッドがかなり大きいことに起因すると考えられる。その影響を除いた新たなスクリプトでは、GPU使用のtorchlibrosaは、同じくGPU使用のtorchaudioと比較して3倍ほど低速という結果が得られた。とはいえ実行時間の桁は揃い、どちらも十分実用に供すると考える。