train.py

import torch
import torch.nn as nn
import pytorch_lightning as pl
import torch.utils.data as data

from utils import voc_single_track
from model import Seq2SeqTranscriber
from dataset import CustomDataset
from utils import token_seg_list_to_midi, unpack_sequence    
from eval import evaluate_midi

class LitTranscriber(pl.LightningModule):
    def __init__(
        self,
        transcriber_args: dict, # Seq2SeqTranscriber モジュールの初期化に必要な引数を含む辞書. 例えばn_melsやn_fft, hop_lengthなど
        lr: float, # 学習率（Learning Rate）。
        lr_decay: float = 1.0, # 学習率の減衰率。デフォルト値は 1.0（減衰なし）。
        lr_decay_interval: int = 1, # 学習率の減衰間隔（エポック単位）。
    ):
        super().__init__()
        self.save_hyperparameters() # 渡された引数をハイパーパラメータとして保存（PyTorch Lightning機能）。
        self.voc_dict = voc_single_track # voc_dict をインスタンス変数に保存。
        self.n_voc = self.voc_dict["n_voc"] # トークンの総数（n_voc）も保存。

        # 渡された引数 transcriber_args を展開して Seq2SeqTranscriber モジュールを初期化。トークン辞書（voc_dict）を追加で渡す。
        self.transcriber = Seq2SeqTranscriber(
            **transcriber_args, voc_dict=self.voc_dict
        )
        # 引数から渡された学習率や減衰関連の設定を保存。
        self.lr = lr
        self.lr_decay = lr_decay
        self.lr_decay_interval = lr_decay_interval

    # 推論時に Seq2SeqTranscriber の推論機能（infer メソッド）を呼び出す関数
    def forward(self, y: torch.Tensor):
        transcriber_infer = self.transcriber.infer(y) # 入力波形（y）を Seq2SeqTranscriber の infer メソッドに渡して推論を実行。
        return transcriber_infer #推論結果として生成されたトークン列を返します。

    # モデルのトレーニング中に1バッチの処理を実行する関数
    def training_step(self, batch, batch_idx):
        y, t = batch # 入力波形（y）と正解トークン列（t）を分割。
        tf_out = self.transcriber(y, t) # 入力データを Seq2SeqTranscriber に渡してtranscriberのforwardで損失値（loss）を取得。出力は、損失値やロジット（各トークンのスコア）を含むオブジェクト。
        loss = tf_out.loss # T5モデルの出力に含まれる損失値（loss）。学習時の損失計算は、CrossEntropyLoss に基づいています。
        t = t.detach() # 正解トークン（t）を計算グラフから切り離して、後続の計算が逆伝播に影響しないようにする。
        mask = t != self.voc_dict["pad"] # 正解トークン列のうち、パディングトークン（pad）を無視するためのマスクを作成。

        # モデルの出力（logits）から最も高いスコアのトークン（argmax(-1)）を取得し、マスクされた正解トークン（t[mask]）と比較。
        # 正しく予測されたトークンの数を、マスクされたトークン全体の数で割って精度を計算。
        accr = (tf_out.logits.argmax(-1)[mask] == t[mask]).sum() / mask.sum()
        #損失値（loss）と精度（accr）を PyTorch Lightning のログ機能を使って記録。
        self.log("train_loss", loss)
        self.log("train_accr", accr)
        return loss # 計算された損失を返します。この値はPyTorch Lightningがバックプロパゲーションを行うために使用します。

    # モデルの検証時に1バッチの処理を実行する関数
    def validation_step(self, batch, batch_idx):
        assert not self.transcriber.training # モデルがトレーニングモードでないことを確認
        y, t = batch
        tf_out = self.transcriber(y, t)
        loss = tf_out.loss
        t = t.detach()
        mask = t != self.voc_dict["pad"]
        accr = (tf_out.logits.argmax(-1)[mask] == t[mask]).sum() / mask.sum()
        self.log("vali_loss", loss)
        self.log("vali_accr", accr)
        return loss

    def test_step(self, batch, batch_idx):
        y, ref_midi = batch # 音声データ（y）と参照MIDI（ref_midi）を分割
        # テストデータセットでは、通常1つのサンプルを1バッチとして処理するため、インデックス0の要素を取り出します。
        y = y[0]
        ref_midi = ref_midi[0]
        with torch.no_grad(): # 推論中は勾配計算が不要なため、この文脈を使ってメモリ消費を削減。
            est_tokens = self.forward(y) # 入力音声 y を用いてモデルに推論を実行し、MIDIトークン列（est_tokens）を生成。モデルの forward メソッドを呼び出して推論を実行。この時forward内では分割されて実行される, 形状： (セグメント数, 最大トークン長)
            unpadded_tokens = unpack_sequence(est_tokens.cpu().numpy()) # 推論結果（est_tokens）からパディングを取り除き、有効なトークン列だけを取得。GPU上のテンソルをCPUに移動してNumPy配列に変換
            unpadded_tokens = [t[1:] for t in unpadded_tokens] # 各トークン列から開始トークン（<sos>など）を除外。
            est_midi = token_seg_list_to_midi(unpadded_tokens) # パディングを除去したトークン列（unpadded_tokens）をMIDI形式に変換。
        dict_eval = evaluate_midi(est_midi, ref_midi) # 参照MIDI（ref_midi）と推論MIDI（est_midi）を比較し、性能を評価。
        # 評価結果をPyTorch Lightningのログ機能で記録。
        dict_log = {}
        for key in dict_eval:
            dict_log["test/" + key] = dict_eval[key]
        self.log_dict(dict_log, batch_size=1)

    def train_dataloader(self):
        # 事前に定義されたデータセットクラス。音声（波形データ）とMIDIファイルをペアとしてロードします。
        dset = CustomDataset(
            split="train",                      # データセットの分割
            sample_rate=16000,
            apply_pedal=True,
            whole_song=False
        )
        # データローダーの作成
        return data.DataLoader(
            dataset=dset, # 使用するデータセットとして、初期化済みの Maestro を指定。
            collate_fn=dset.collate_batch, # バッチ処理時のデータ整形関数。
            batch_size=64, # 1バッチに含めるサンプル数。
            shuffle=True, # データセットをランダムにシャッフル。
            pin_memory=True, # データをピン留め（ページロック）してメモリに保存。
            num_workers=32, # データローディングに使用するプロセス数。
        )

    def val_dataloader(self):
        # 事前に定義されたデータセットクラス。音声（波形データ）とMIDIファイルをペアとしてロードします。
        dset = CustomDataset(
            split="valid",                      # データセットの分割
            sample_rate=16000,
            apply_pedal=True,
            whole_song=False
        )
        # データローダーの作成
        return data.DataLoader(
            dataset=dset, # 使用するデータセットとして、初期化済みの Maestro を指定。
            collate_fn=dset.collate_batch, # バッチ処理時のデータ整形関数。
            batch_size=32, # 1バッチに含めるサンプル数。
            shuffle=True, # データセットをランダムにシャッフル。
            pin_memory=True, # データをピン留め（ページロック）してメモリに保存。
            num_workers=16, # データローディングに使用するプロセス数。
        )
    def test_dataloader(self):
        dset = CustomDataset(
            split="test",                      # データセットの分割
            sample_rate=16000,
            apply_pedal=True,
            whole_song=True
        )
        return data.DataLoader(
            dataset=dset,
            collate_fn=dset.collate_wholesong,
            batch_size=1,
            shuffle=False, # テストデータの順序を固定。理由: 再現性を確保するため。
            pin_memory=True,
        )

    # モデルの学習時に、損失関数の値を最小化するようにパラメータを更新するオプティマイザ（最適化手法）を指定する関数
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=self.lr)
    # Adamオプティマイザ の改良版であり、L2正則化の代わりに ウェイトデカイ（Weight Decay） を導入したオプティマイザ。
    # self.parameters(): モデルの学習可能なすべてのパラメータ（重みとバイアス）を取得。

from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger

# チェックポイント保存の設定
if __name__ == "__main__":
    checkpoint_callback = ModelCheckpoint(
        dirpath="/content/drive/MyDrive/B4/Humtrans/m2m100/checkpoints",  # 保存先
        filename="epoch{epoch:02d}-vali_loss{vali_loss:.2f}",  # ファイル名フォーマット
        save_top_k=3,  # 上位3つのチェックポイントを保存
        monitor="vali_loss",  # 訓練データの損失を監視
        mode="min",  # 損失が小さいほど良いモデル
        every_n_epochs=1,  # 1エポックごとに保存
    )

    # ロガー設定
    logger = TensorBoardLogger(
        save_dir="/content/drive/MyDrive/B4/Humtrans/m2m100/logs",
        name="transcription"
    )

    # トレーナーの設定
    trainer = pl.Trainer(
        logger=logger,
        callbacks=[checkpoint_callback],  # チェックポイントコールバックを追加
        enable_checkpointing=True,
        accelerator="gpu",
        devices=1,
        max_epochs=400,
    )

    # モデルの引数
    args = {
        "n_mels": 128,
        "sample_rate": 16000,
        "n_fft": 1024,
        "hop_length": 128,
    }

    lightning_module = LitTranscriber(
        transcriber_args=args,
        lr=1e-4, # 0.0001
        lr_decay=0.99, # 各エポックの度に0.99倍して減衰させる
    )

    # 学習の開始
    trainer.fit(lightning_module)