PyTorch 2.0 チュートリアル : 入門 : モデルのセーブ & ロード (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/19/2023 (2.0.0)

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

• Introduction to PyTorch : Learn the Basics : Save & Load the Model

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

PyTorch 2.0 チュートリアル : 入門 : モデルのセーブ & ロード

このセクションではセーブ、ロードそしてモデル予測を実行することでモデル状態をどのように永続化するかを見ます。

import torch
import torchvision.models as models

モデル重みをセーブしてロードする

PyTorch モデルは学習されたパラメータを state_dict と呼ばれる、内部状態辞書にストアします。これらは torch.save メソッドを通して永続化できます :

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'model_weights.pth')

/opt/conda/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning:

The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.

/opt/conda/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning:

Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.

Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to /var/lib/jenkins/.cache/torch/hub/checkpoints/vgg16-397923af.pth

  0%|          | 0.00/528M [00:00<?, ?B/s]
  0%|          | 1.99M/528M [00:00<00:26, 20.7MB/s]
  1%|1         | 5.68M/528M [00:00<00:18, 28.8MB/s]
  2%|1         | 10.4M/528M [00:00<00:14, 36.9MB/s]
  3%|2         | 14.7M/528M [00:00<00:13, 40.0MB/s]
  4%|3         | 19.0M/528M [00:00<00:12, 41.7MB/s]
  4%|4         | 23.0M/528M [00:00<00:12, 41.5MB/s]
  5%|5         | 28.1M/528M [00:00<00:11, 45.2MB/s]
  7%|6         | 34.3M/528M [00:00<00:10, 51.6MB/s]
  8%|7         | 39.8M/528M [00:00<00:09, 53.2MB/s]
  8%|8         | 44.8M/528M [00:01<00:10, 49.2MB/s]
 10%|9         | 51.4M/528M [00:01<00:09, 54.9MB/s]
 11%|#         | 56.7M/528M [00:01<00:10, 47.5MB/s]
 12%|#1        | 61.5M/528M [00:01<00:12, 38.8MB/s]
 12%|#2        | 65.6M/528M [00:01<00:12, 39.3MB/s]
 14%|#3        | 71.4M/528M [00:01<00:10, 44.6MB/s]
 14%|#4        | 76.0M/528M [00:01<00:11, 41.8MB/s]
 15%|#5        | 81.0M/528M [00:01<00:10, 44.6MB/s]
 16%|#6        | 86.5M/528M [00:02<00:09, 47.9MB/s]
 17%|#7        | 91.2M/528M [00:02<00:10, 42.2MB/s]
 18%|#8        | 95.7M/528M [00:02<00:10, 43.3MB/s]
 19%|#9        | 101M/528M [00:02<00:09, 45.4MB/s]
 20%|#9        | 105M/528M [00:02<00:10, 44.2MB/s]
...
 90%|########9 | 474M/528M [00:11<00:01, 34.2MB/s]
 91%|######### | 478M/528M [00:11<00:01, 34.6MB/s]
 91%|#########1| 483M/528M [00:11<00:01, 39.1MB/s]
 92%|#########2| 487M/528M [00:12<00:01, 33.9MB/s]
 93%|#########2| 490M/528M [00:12<00:01, 30.3MB/s]
 94%|#########3| 496M/528M [00:12<00:00, 34.2MB/s]
 95%|#########4| 500M/528M [00:12<00:00, 35.4MB/s]
 96%|#########5| 505M/528M [00:12<00:00, 41.3MB/s]
 97%|#########6| 512M/528M [00:12<00:00, 48.2MB/s]
 98%|#########7| 516M/528M [00:12<00:00, 46.8MB/s]
 99%|#########8| 521M/528M [00:12<00:00, 46.6MB/s]
100%|#########9| 526M/528M [00:13<00:00, 13.3MB/s]
100%|##########| 528M/528M [00:14<00:00, 39.1MB/s]

モデル重みをロードするには、最初に同じモデルのインスタンスを作成してから load_state_dict() メソッドを使用してパラメータをロードする必要があります。

model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

Note

ドロップアウトとバッチ正規化層を評価モードに設定するために推論の前に model.eval() メソッドを確実に呼び出してください。これを行なわなければ一貫性のない推論結果を生成します。

形状 (Shape) と共にモデルをセーブしてロードする

モデル重みをロードするとき、モデルクラスを最初にインスタンス化することが必要でした、何故ならばクラスがネットワークの構造を定義するからです。私たちはこのクラスの構造をモデルと一緒にセーブすることを望むかもしれません、その場合モデル (そして model.state_dict() ではない) をセービング関数に渡すことができます :

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

そしてこのようにモデルをロードできます :

model = torch.load('model.pth')

Note

このアプローチはモデルをシリアライズするとき Python pickle モジュールを使用しますので、モデルをロードするときそれは実際のクラス定義が利用可能であることに依存します。

関連チュートリアル

• Saving and Loading a General Checkpoint in PyTorch

以上