MusicFMを動かす

概要

MusicFMというモデルを動かしてみました。こちらは A Foundation Model for Music Informatics という論文のモデルです。
https://arxiv.org/abs/2311.03318

こちらのモデルは音楽から情報を取得するときの基盤モデルという位置付けのようです。この分野ではMERTというモデルが性能が高く、よく使われていますが、MusicFMはその後発の研究です。githubのスターの数が多いので注目されているようです。

実装

こちらでモデルの実装が公開されています。ライセンスはMITです。学習済みモデルのライセンスは特に記述がないので、同じMITライセンスと判断してよさそうです。
https://github.com/minzwon/musicfm

学習

MusicFMの学習にはFMAというデータセットが使われています。こちらもMITライセンスです。FMAデータセットはsmall, medium, large, fullと4種類ありますが、MusicFMではlargeを使って学習しています。シンプルにzipファイルがダウンロードできるのは便利ですね。
https://github.com/mdeff/fma

論文中には記述がありませんがgithubではMillion Song Dataset(MSD)で学習したモデルも公開されていました。こちらのほうが性能が高いようです。
http://millionsongdataset.com/

動かしてみる

環境設定

以下の環境で動作確認できました。
Python3.11
Ubuntu22.04
CUDA11.8
RTX4090

pytorch-triton                3.0.0+901819d2b6
torch                         2.3.0.dev20240301+cu118
torchaudio                    2.2.0.dev20240301+cu118
torchvision                   0.18.0.dev20240301+cu118
einops                        0.7.0

埋め込みを取得

ソースを取得

git clone https://github.com/minzwon/musicfm.git

学習済みモデルを取得。FMAのほうをダウンロードします。

wget -P musicfm/data/ https://huggingface.co/minzwon/MusicFM/resolve/main/fma_stats.json
wget -P musicfm/data/ https://huggingface.co/minzwon/MusicFM/resolve/main/pretrained_fma.pt

埋め込みを取得。wavファイルは適当な乱数配列を使っているようです。githubの記述ではMSDになっているので、FMAに書き換えています。

HOME_PATH = "./" # path where you cloned musicfm

import os
import sys
import torch

sys.path.append(HOME_PATH)
from musicfm.model.musicfm_25hz import MusicFM25Hz

# dummy audio (30 seconds, 24kHz)
wav = (torch.rand(4, 24000 * 30) - 0.5) * 2

# load MusicFM
musicfm = MusicFM25Hz(
    is_flash=False,
    stat_path=os.path.join(HOME_PATH, "musicfm", "data", "fma_stats.json"),
    model_path=os.path.join(HOME_PATH, "musicfm", "data", "pretrained_fma.pt"),
)

# to GPUs
wav = wav.cuda()
musicfm = musicfm.cuda()

# get embeddings
musicfm.eval()
emb = musicfm.get_latent(wav, layer_ix=7)

print(emb)

実行結果

tensor([[[-4.3217e-01,  1.1156e+00,  7.2494e-01,  ..., -1.5156e+00,
           5.8803e-01, -1.3097e-01],
         [ 1.8403e-01,  6.7820e-01,  1.0176e+00,  ..., -1.4453e+00,
           4.6587e-01, -1.2313e+00],
         [-1.4451e-01,  2.9229e-01,  1.0401e+00,  ..., -9.5231e-01,
           8.3809e-01, -1.3208e+00],
         ...,