PyTorch モデルの保存と読み込み

ディープラーニングプロジェクトにおいて、モデルの保存と読み込みは非常に重要な要素であり、主な理由は次のとおりです：

  1. トレーニング中断からの復元：トレーニングプロセスが予期せず中断された場合、保存ポイントからトレーニングを継続できます
  2. モデルデプロイ：学習済みモデルを本番環境にデプロイする
  3. モデル共有：チームメンバー間でモデルの成果物を簡単に共有できる
  4. 転移学習：事前学習済みモデルを他のタスク用に保存する
  5. 性能評価：異なるトレーニング段階のモデルを保存して比較する

基本的な保存と読み込み方法

モデル全体を保存

これは最も簡単な方法で、モデルのアーキテクチャとパラメータを保存します：

例

import torch
import torchvision.models as models

# モデルを作成してトレーニングする
model = models.resnet18(pretrained=True)
# ... トレーニングコード ...

# モデル全体を保存
torch.save(model, 'model.pth')

# モデル全体を読み込み
loaded_model = torch.load('model.pth')

利点：

  • コードがシンプルで直感的
  • 完全なモデル構造を保存する

欠点：

  • ファイルサイズが大きい
  • モデルクラスの定義に依存する

モデルパラメータのみ保存（推奨方法）

より推奨される方法は、モデルの状態辞書（state_dict）のみを保存することです：

例

# モデルパラメータを保存
torch.save(model.state_dict(), 'model_weights.pth')

# モデルパラメータを読み込み
model = models.resnet18()  # 最初に同じアーキテクチャのモデルを作成する必要があります
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # 評価モードに設定

利点：

  • ファイルがより小さい
  • より柔軟で、異なるアーキテクチャに読み込める
  • 互換性がより良い

トレーニング状態の保存と読み込み

実際のプロジェクトでは、通常オプティマイザの状態やエポックなどの情報も保存する必要があります：

例

# チェックポイントを保存
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    # 保存が必要な他の情報を追加できます
}

torch.save(checkpoint, 'checkpoint.pth')

# チェックポイントを読み込み
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()  # または model.train() を必要に応じて使用

デバイスをまたいだモデルの読み込み

CPU/GPU互換性の処理

例

# 保存時にmap_locationを指定
torch.save(model.state_dict(), 'model_weights.pth')

# CPUに読み込む（モデルがGPUでトレーニングされた場合）
device = torch.device('cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))

# GPUに読み込む
device = torch.device('cuda')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
model.to(device)

マルチGPUトレーニングモデルの読み込み

例

# マルチGPUモデルを保存
torch.save(model.module.state_dict(), 'multigpu_model.pth')

# シングルGPUに読み込み
model = ModelClass()
model.load_state_dict(torch.load('multigpu_model.pth'))

モデル変換と互換性

PyTorchバージョンの互換性

例

# 保存時に_use_new_zipfile_serialization=Trueを指定して、より良い互換性を得る
torch.save(model.state_dict(), 'model.pth', _use_new_zipfile_serialization=True)

TorchScriptへの変換

例

# モデルをTorchScript形式に変換
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'model_scripted.pt')

# TorchScriptモデルを読み込み
loaded_script = torch.jit.load('model_scripted.pt')

ベストプラクティスとよくある問題

ベストプラクティス

  1. 命名規則：意味のあるファイル名を使用する。例えばresnet18_epoch50.pth
  2. 定期的な保存：数エポックごとにチェックポイントを保存する
  3. 読み込みの検証：保存後すぐに読み込み機能をテストする
  4. ドキュメント記録：モデルアーキテクチャとトレーニングパラメータを記録する
  5. バージョン管理：モデルファイルをバージョン管理システムに含める

よくある問題の解決方法

問題1：Missing key(s) in state_dict
解決：モデルアーキテクチャが完全に一致することを確認するか、次を使用しますstrict=Falseパラメータ：

model.load_state_dict(torch.load('model.pth'), strict=False)

問題2：CUDA out of memory
解決：読み込み時はまずCPUに配置します：

例

model.load_state_dict(torch.load('model.pth', map_location='cpu'))

問題3：古いバージョンのモデルを読み込めない
解決：異なるPyTorchバージョンで読み込むか、モデル形式を変換してみてください

実際の応用例

画像分類モデルの保存と読み込みの流れ

完全なコード例

例

import torch
import torch.nn as nn
import torch.optim as optim

# シンプルなモデルを定義
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)
   
    def forward(self, x):
        return self.fc(x)

# 初期化
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# トレーニングプロセスをシミュレーション
for epoch in range(5):
    # トレーニングステップをシミュレーション
    inputs = torch.randn(32, 10)
    labels = torch.randint(0, 2, (32,))
   
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
   
    # 2エポックごとにチェックポイントを保存
    if epoch % 2 == 0:
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item(),
        }
        torch.save(checkpoint, f'checkpoint_epoch{epoch}.pth')
        print(f'Checkpoint saved at epoch {epoch}')

# 最終保存
torch.save(model.state_dict(), 'final_model.pth')

# 読み込み例
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('final_model.pth'))
loaded_model.eval()

# 読み込んだモデルをテスト
test_input = torch.randn(1, 10)
with torch.no_grad():
    output = loaded_model(test_input)
print(f'Test output: {output}')
その他の拡張