PyTorchの「torch.save」関数：モデルの保存と復元を分かりやすく解説

PyTorchにおける「torch.save」：モデルの保存と復元

使い方

torch.save(model, PATH)

このコードは、model というモデルオブジェクトを PATH というファイルパスに保存します。モデルには、モデルのアーキテクチャ、学習済みのパラメータ、およびその他の必要な状態情報が含まれます。

例

import torch

# モデルを定義して訓練する
model = ...
model.train()
# ...

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

この例では、model というモデルオブジェクトを my_model.pth というファイルに保存します。

モデルの復元

保存済みのモデルをロードするには、torch.load 関数を使用します。

model = torch.load(PATH)

このコードは、PATH ファイルからモデルオブジェクトを読み込みます。

例

import torch

# モデルを復元
model = torch.load('my_model.pth')

# モデルを使用して推論を行う
# ...

この例では、my_model.pth ファイルからモデルオブジェクトをロードし、推論に使用します。

注意事項

• torch.save は、CPU 上のモデルのみを保存できます。GPU 上のモデルを保存するには、まず .cpu() メソッドを使用して CPU にコピーする必要があります。
• モデルを保存する前に、モデルの訓練モードをオフにする必要があります。これは、model.eval() メソッドを使用して行うことができます。
• モデルをロードするときは、モデルのアーキテクチャと訓練済みのパラメータが保存されたファイルと同じであることを確認する必要があります。

補足

• torch.save は、モデルだけでなく、任意の Python オブジェクトを保存するのにも使用できます。
• モデルをより効率的に保存するには、torch.jit.save 関数を使用することができます。この関数は、モデルを TorchScript モジュールに変換してから保存します。TorchScript モジュールは、ネイティブコードに変換できるため、より速く実行できます。

この説明が、PyTorchにおける "torch.save" の理解に役立つことを願っています。ご不明な点がございましたら、お気軽にご連絡ください。

PyTorch torch.save のサンプルコード

モデルの保存

import torch

# モデルを定義して訓練する
model = ...
model.train()
# ...

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

このコードは、model というモデルオブジェクトを my_model.pth というファイルに保存します。

モデルとオプティマイザの保存

import torch

# モデルとオプティマイザを定義して訓練する
model = ...
optimizer = ...
# ...

# モデルとオプティマイザを保存
checkpoint = {'model': model, 'optimizer': optimizer}
torch.save(checkpoint, 'checkpoint.pth')

このコードは、model というモデルオブジェクトと optimizer というオプティマイザオブジェクトを checkpoint.pth というファイルに保存します。

カスタム状態の保存

import torch

# モデルを定義して訓練する
model = ...
model.train()
# ...

# カスタム状態を定義
custom_state = {'data': some_data, 'running_loss': 0.0}

# モデルとカスタム状態を保存
checkpoint = {'model': model, 'custom_state': custom_state}
torch.save(checkpoint, 'checkpoint.pth')

このコードは、model というモデルオブジェクトと custom_state というカスタム状態辞書を checkpoint.pth というファイルに保存します。

モデルのロード

import torch

# モデルをロード
model = torch.load('my_model.pth')

# モデルを使用して推論を行う
# ...

このコードは、my_model.pth ファイルからモデルオブジェクトをロードし、推論に使用します。

モデルとオプティマイザのロード

import torch

# モデルとオプティマイザをロード
checkpoint = torch.load('checkpoint.pth')
model = checkpoint['model']
optimizer = checkpoint['optimizer']

# トレーニングを再開する
# ...

このコードは、checkpoint.pth ファイルからモデルオブジェクトとオプティマイザオブジェクトをロードし、トレーニングを再開します。

カスタム状態のロード

import torch

# モデルとカスタム状態をロード
checkpoint = torch.load('checkpoint.pth')
model = checkpoint['model']
custom_state = checkpoint['custom_state']

# カスタム状態を使用する
data = custom_state['data']
running_loss = custom_state['running_loss']
# ...

このコードは、checkpoint.pth ファイルからモデルオブジェクトとカスタム状態辞書をロードし、カスタム状態を使用します。

これらのサンプルコードは、torch.save 関数の基本的な使用方法を示しています。詳細については、PyTorch のドキュメントを参照してください。

何かご質問があれば、気軽にお尋ねください。

1. torch.save を使用する: これは、PyTorchモデルを保存する最も一般的で簡単な方法です。この関数は、モデルとその状態を単一のファイルに保存します。

torch.save(model, PATH)

1. TorchScript を使用する: TorchScriptは、PyTorchモデルをPython非依存の形式に変換するためのツールです。TorchScriptモデルは、ネイティブコードに変換できるため、より速く実行できます。TorchScriptモデルを保存するには、torch.jit.save 関数を使用します。

torch.jit.save(model, PATH)

一般的に、以下の場合は torch.save を使用する方が良いでしょう。

• モデルを後で別の環境でロードしたい場合
• モデルをデバッグする必要がある場合
• モデルをカスタムトレーニングしたい場合

以下の場合は TorchScript を使用する方が良いでしょう。

• モデルを可能な限り高速に実行したい場合
• モデルを軽量化したい場合
• モデルをWebブラウザやモバイルアプリでデプロイしたい場合

その他の保存方法

上記以外にも、PyTorchモデルを保存する方法はいくつかあります。

• ONNX を使用する: ONNXは、機械学習モデルの交換可能な形式です。PyTorchモデルをONNX形式で保存するには、torch.onnx.export 関数を使用します。
• Protobuf を使用する: Protobufは、構造化データを効率的にシリアル化するためのバイナリ形式です。PyTorchモデルをProtobuf形式で保存するには、モデルをカスタムメッセージとして定義する必要があります。

何かご質問があれば、気軽にお尋ねください。