実務でよく使うPyTorchの機能まとめ③｜modelの保存・読み込み(torch.save・torch.load)

基本的な用法

torch.save

PyTorchで学習させたmodelの保存にはtorch.saveを用います。

import torch
# https://pytorch.org/docs/main/generated/torch.save.html

# Save to file
x = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
torch.save(x, "tensor.pt")

上記を実行するとカレントディレクトリにtensor.ptというファイルが保存されます。

torch.load

保存させたmodelの読み込みにはtorch.loadを用います。下記が基本的な実行例です。

model = torch.load("tensor.pt", weights_only=True)
print(type(model), model)

・実行結果

<class 'torch.Tensor'> tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

CPU環境で呼び出す場合はtorch.loadの引数にmap_location=torch.device("cpu")を入れて実行します。

抑えておくと良い使い方