Griffin-Limアルゴリズムの実行時間をlibrosaとtorchaudioで比較してみた話

はじめに

音声の振幅スペクトルから位相を推定し、元の音声を復元するためのGriffin-Limアルゴリズムが知られている。
Griffin-Limアルゴリズムはlibrosaパッケージとtorchaudioパッケージの両方に実装されている。

• librosa

librosa.org

• torchaudio

pytorch.org

今回はGriffin-Limアルゴリズムに関して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

スクリプト

"""
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 numpy as np
import torch
import torchaudio
import torchaudio.transforms as T

N_FFT = 1024
WIN_LENGTH = None
HOP_LENGTH = 512
N_ITER = 50

if __name__ == "__main__":

    waveform, sample_rate = librosa.load(librosa.ex("trumpet"))

    # librosa
    start = time.time()
    specgram = np.abs(
        librosa.stft(
            waveform, n_fft=N_FFT, hop_length=HOP_LENGTH, win_length=WIN_LENGTH
        )
    )
    y_inv = librosa.griffinlim(
        specgram, n_iter=N_ITER, hop_length=HOP_LENGTH, win_length=WIN_LENGTH
    )
    elapsed_time = time.time() - start
    print("elapsed_time (librosa): {0:.6f}".format(elapsed_time) + "[sec]")

    # load waveform as torch.Tensor
    waveform, sample_rate = torchaudio.load(librosa.ex("trumpet"))
    waveform_cuda = waveform.cuda()

    # torchaudio (CPU)
    specgram = T.Spectrogram(
        n_fft=N_FFT,
        win_length=WIN_LENGTH,
        hop_length=HOP_LENGTH,
    )
    griffin_lim = T.GriffinLim(
        n_fft=N_FFT,
        n_iter=N_ITER,
        win_length=WIN_LENGTH,
        hop_length=HOP_LENGTH,
    )
    start = time.time()
    reconstructed = griffin_lim(specgram(waveform))
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; CPU): {0:.6f}".format(elapsed_time) + "[sec]")

    # torchaudio (GPU)
    specgram = T.Spectrogram(
        n_fft=N_FFT,
        win_length=WIN_LENGTH,
        hop_length=HOP_LENGTH,
    ).cuda()
    griffin_lim = T.GriffinLim(
        n_fft=N_FFT,
        n_iter=N_ITER,
        win_length=WIN_LENGTH,
        hop_length=HOP_LENGTH,
    ).cuda()
    torch.cuda.synchronize()
    start = time.time()
    reconstructed = griffin_lim(specgram(waveform_cuda))
    torch.cuda.synchronize()
    elapsed_time = time.time() - start
    print("elapsed_time (torchaudio; GPU): {0:.6f}".format(elapsed_time) + "[sec]")

実行結果

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

elapsed_time (librosa): 0.264543[sec]
elapsed_time (torchaudio; CPU): 0.147203[sec]
elapsed_time (torchaudio; GPU): 0.022305[sec]

まとめ

位相復元に関しても、torchaudioのCPU版を使うだけで実行時間の短縮が望めそうである。GPU版も悪くない結果である。今回は内部的に実装されているアルゴリズムを考慮したフェアな比較というよりは、いちユーザが関数を外から触るときの体感を重視した。