piperモデルからつくよみちゃんデータセットを使って追加学習を行う

初めに

以下の記事でjvsデータセットを用いてpiperモデルの日本語化対応を行いました。

ayousanz.hatenadiary.jp

しかし、このモデルはあくまで事前学習モデルであって特定の話者に特化しているものではありません。そのため、今回は 以下のつくよみちゃんデータセットを用いて 追加学習を行い 特定話者の声のモデルを作成していきます

tyc.rei-yumesaki.net

Demo

追加学習をした後の音声は以下のようになり、つくよみちゃんの声になっていることがわかります

youtu.be

開発環境

前処理

追加学習をする際に使用するデータセットを ljspeechフォーマットに変換しておく必要があります。

事前学習モデルから 追加学習用のモデルを作成

事前学習モデルからそのまま別のデータセットを使って追加学習をすることができないため、以下の対応を事前に行います。

  • モデルの重みだけをロードし、オプティマイザの状態はロードせずに、新しく初期化
  • 話者埋め込み層を削除
import torch
import sys
import os

def create_partial_checkpoint_for_finetuning(original_ckpt_path, new_ckpt_path):
    """
    ファインチューニングのために、不整合なレイヤー(話者関連の層)と
    オプティマイザの状態を削除した新しいチェックポイントを作成します。
    """
    print(f"Loading original checkpoint from: {original_ckpt_path}")
    
    try:
        # CPUにロードしてGPUメモリを節約
        checkpoint = torch.load(original_ckpt_path, map_location="cpu")
    except FileNotFoundError:
        print(f"ERROR: Checkpoint file not found at {original_ckpt_path}", file=sys.stderr)
        return

    original_state_dict = checkpoint["state_dict"]
    
    # 削除するキーのリスト (エラーログに表示されたUnexpected keys)
    # これらは主に、話者数が100から1に変わったことで不要になったレイヤーです。
    keys_to_remove = [
        # 話者埋め込み層と関連する条件付けレイヤー
        "model_g.emb_g.weight",
        "model_g.dec.cond.weight",
        "model_g.dec.cond.bias",
        "model_g.enc_q.enc.cond_layer.bias",
        "model_g.enc_q.enc.cond_layer.weight_g",
        "model_g.enc_q.enc.cond_layer.weight_v",
        "model_g.dp.cond.weight",
        "model_g.dp.cond.bias",
    ]
    # フローの条件付けレイヤーも動的にリストに追加
    for i in [0, 2, 4, 6]:
        for suffix in ["bias", "weight_g", "weight_v"]:
            keys_to_remove.append(f"model_g.flow.flows.{i}.enc.cond_layer.{suffix}")
    
    keys_to_remove_set = set(keys_to_remove)

    # 新しいstate_dictから、削除対象のキーを除外して作成
    new_state_dict = {key: value for key, value in original_state_dict.items() if key not in keys_to_remove_set}

    # チェックポイントのstate_dictを新しいものに更新
    checkpoint["state_dict"] = new_state_dict
    
    # オプティマイザの状態をチェックポイントから削除 (★今回の修正点★)
    if "optimizer_states" in checkpoint:
        del checkpoint["optimizer_states"]
        print("Removed optimizer states from the checkpoint.")
    
    # 新しい部分的なチェックポイントを保存
    torch.save(checkpoint, new_ckpt_path)

    print("\n--- Checkpoint Modification Summary ---")
    print(f"Original state_dict had {len(original_state_dict)} keys.")
    print(f"Removed {len(original_state_dict) - len(new_state_dict)} keys.")
    print(f"New state_dict has {len(new_state_dict)} keys.")
    print("---------------------------------------")
    print(f"\nPartial checkpoint for fine-tuning saved to: {new_ckpt_path}")
    print("\nNow, use this new partial checkpoint path for the --resume_from_checkpoint argument.")

if __name__ == "__main__":
    # --- ★設定箇所★ ---
    # 元のJVSモデルのチェックポイントパス
    original_checkpoint = "/data/piper_jvs_preprocessed_espeak_ms/lightning_logs/version_0/checkpoints/epoch=579-step=129920.ckpt"

    # 保存する新しい「部分的チェックポイント」のパスとファイル名
    partial_checkpoint = "/data/piper_jvs_preprocessed_espeak_ms/lightning_logs/version_0/checkpoints/epoch=579_partial_for_finetune.ckpt"
    # --- 設定ここまで ---

    create_partial_checkpoint_for_finetuning(original_checkpoint, partial_checkpoint)

実行すると以下のようなログが出てきます。こちらで追加学習する用のモデルに変換できました

--- Checkpoint Modification Summary ---
Original state_dict had 804 keys.
Removed 20 keys.
New state_dict has 784 keys.
---------------------------------------

追加学習

以下で追加学習に使用する先ほど変換したモデルとデータセットのパスを環境変数に定義して、学習を行なっていきます

export PARTIAL_CHECKPOINT_PATH="/data/piper_jvs_preprocessed_espeak_ms/lightning_logs/version_0/checkpoints/epoch=579_partial_for_finetune.ckpt"
export PREPROCESSED_TSUKUYOMI_DIR="/data/piper_tsukuyomi_preprocessed_rawtext"

python3 -m piper_train \
  --dataset-dir ${PREPROCESSED_TSUKUYOMI_DIR} \
  --accelerator 'gpu' \
  --devices 1 \
  --batch-size 16 \
  --max_epochs 300 \
  --checkpoint-epochs 5 \
  --precision 32 \
  --quality medium \
  --resume_from_checkpoint "${PARTIAL_CHECKPOINT_PATH}"

この時のepochs数は、事前学習時のepoch数 + 追加学習したいepoch数の合計になります

onnxに変換

いつもと同じく変換をします

python3 -m piper_train.export_onnx \
  "${FINETUNED_CKPT_PATH}" \
  "${EXPORT_DIR}/${MODEL_NAME}.onnx"

推論

以下で推論を行います

echo "${TEST_SENTENCE}" | \
  piper \
    -m "${EXPORT_DIR}/${MODEL_NAME}.onnx" \
    --output_file "${OUTPUT_WAV_PATH}"