PyTorch 1.8 チュートリアル : レシピ : 基本 :- マルチモデルを一つのファイルでセーブ/ロードする (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 04/29/2021 (1.8.1+cu102)

* 本ページは、PyTorch 1.8 Tutorials の以下のページを翻訳した上で適宜、補足説明したものです：

• PyTorch Recipes : Basics : Saving and Loading Multiple Models in One File using PyTorch

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

レシピ : 基本 :- マルチモデルを一つのファイルでセーブ/ロードする

マルチモデルをセーブしてロードすることは以前に訓練したモデルを再利用するために有用であり得ます。

イントロダクション

GAN、sequence-to-sequence モデルあるいはモデルのアンサンブルのような、マルチ torch.nn.Module から成るモデルをセーブするとき、各モデルの state_dict と対応する optimizer の辞書をセーブしなければなりません。訓練を再開することを補助するかもしれない任意の他の項目もそれらを単純に辞書に追加することによりセーブできます。モデルをロードするには、最初にモデルと optimizer を初期化してから、torch.load() を使用してローカルで辞書をロードします。ここから、想定する辞書に単純に問い合せることでセーブされた項目に容易にアクセスできます。このレシピでは、PyTorch を使用してマルチモデルを一つのファイルにどのようにセーブするか実演します。

セットアップ

始める前に、torch をそれがまだ利用可能でないならばインストールする必要があります。

pip install torch

ステップ

1. データをロードするために総ての必要なライブラリをインポートする。
2. ニューラルネットワークを定義して初期化する。
3. optimizer を初期化する。
4. マルチモデルをセーブする。
5. マルチモデルをロードする。

1. データをロードするために必要なライブラリをインポートする

このレシピのため、torch とその補助 torch.nn と torch.optim を使用します。

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

2. ニューラルネットワークを定義して初期化する

例のために、訓練画像のためのニューラルネットワークを作成します。更に学習するためには ニューラルネットワークを定義する レシピを見てください。モデルを最終的にセーブするために 2 つの変数を構築します。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

netA = Net()
netB = Net()

3. optimizer を初期化する

作成した各モデルのための optimizer を構築するために SGD with モメンタムを使用します。

optimizerA = optim.SGD(netA.parameters(), lr=0.001, momentum=0.9)
optimizerB = optim.SGD(netB.parameters(), lr=0.001, momentum=0.9)

4. マルチモデルをセーブする

総ての関連情報を集めて辞書を構築します。

# Specify a path to save to
PATH = "model.pt"

torch.save({
            'modelA_state_dict': netA.state_dict(),
            'modelB_state_dict': netB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            }, PATH)

5. マルチモデルをロードする

最初にモデルと optimizer を初期化してから、辞書をローカルでロードすることを忘れないでください。

modelA = Net()
modelB = Net()
optimModelA = optim.SGD(modelA.parameters(), lr=0.001, momentum=0.9)
optimModelB = optim.SGD(modelB.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

推論を実行する前に dropout とバッチ正規化層を評価モードに設定するために model.eval() を呼び出さなければなりません。これをし損なうと一貫性のない推論結果をもたらします。

訓練を再開することを望む場合には、これらの層が訓練モードにあることを確かなものにするために model.train() を呼び出してください。

Congratulations! PyTorch でマルチモデルを成功的にセーブしてロードしました。

以上