PyTorch torch.save と torch.load 関数

Pytorch torch 参考手册 Pytorch torch リファレンスマニュアル

torch.saveとtorch.loadこれらは、PyTorch でテンソル、モデル、その他の Python オブジェクトをシリアライズ（保存）およびデシリアライズ（ロード）するための関数です。

これらの関数は、学習済みモデルの保存や、トレーニングを再開するためのチェックポイントの保存などのシナリオで不可欠です。

関数定義

torch.save(obj, f, pickle_module, pickle_protocol)
torch.load(f, map_location, pickle_module, weights_only)

torch.save の引数:

obj: 保存するオブジェクト。テンソル、モデル、辞書、リストなどを指定できます。
f: ファイルパス（文字列またはファイルオブジェクト）。
pickle_module(オプション): シリアライズに使用するモジュール。
pickle_protocol(オプション): シリアライズプロトコルのバージョン。

torch.load の引数:

f: ファイルパス（文字列またはファイルオブジェクト）。
map_location(オプション): ストレージを異なるデバイスにマッピングする方法を指定します。
pickle_module(オプション): デシリアライズに使用するモジュール。
weights_only(bool, オプション): Python オブジェクトをロードせずに重みのみをロードするかどうか。

使用例

例1: テンソルの保存とロード

例

import torch

# テンソルをいくつか作成
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)

# テンソルをファイルに保存
torch.save({'x': x, 'y': y}, 'tensors.pth')

# ファイルからテンソルをロード
loaded = torch.load('tensors.pth')

print("ロードしたデータ:", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])

出力結果は次のとおりです：

加载的数据: {'x': tensor([1, 2, 3, 4, 5]), 'y': tensor([[-0.3042, -0.9077, -1.0826,  0.9333],
        [ 0.0551,  0.6728,  0.5942, -0.1522],
        [-0.3744,  0.9239, -0.2104, -0.5239]])}
x: tensor([1, 2, 3, 4, 5])
y: tensor([[-0.3042, -0.9077, -1.0826,  0.9333],
        [ 0.0551,  0.6728,  0.5942, -0.1522],
        [-0.3744,  0.9239, -0.2104, -0.5239]])

例2: モデルの保存とロード

例

import torch
import torch.nn as nn

# シンプルなニューラルネットワークを定義する
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)

def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

# モデルのインスタンスを作成する
model = SimpleNet()

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

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

print("モデルが保存され、ロードされました")
print(loaded_model)

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

例

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)

def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

model = SimpleNet()

# モデルパラメータのみを保存（推奨方法）
torch.save(model.state_dict(), 'model_weights.pth')

# 新しいモデルを作成し、パラメータをロードする
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))

print("モデルパラメータが保存され、ロードされました")
print(new_model.state_dict().keys())

出力結果は次のとおりです：

模型参数已保存并加载 odict_keys(['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'])

モデルパラメータのみを保存することをお勧めします（state_dict）、モデル全体を保存するのではなく。これにより、異なるアーキテクチャのモデル間でパラメータを再利用できます。

例4: CPU と GPU 間の移行

例

import torch

# GPU 上でモデルを保存したと仮定します
if torch.cuda.is_available():
x = torch.randn(2, 3, device='cuda')
torch.save(x, 'tensor_gpu.pth')

# CPU 上でロード
x_cpu = torch.load('tensor_gpu.pth', map_location='cpu')
print("CPU にロード:", x_cpu.device)

使用map_locationパラメータを使用すると、テンソルを異なるデバイスにロードできます。

注意事項

保存ファイルには通常、.pth 或 .pt拡張子を使用します。
モデルパラメータのみを保存することをお勧めします（state_dict）、モデル全体を保存するのではなく。
ロード時には、セキュリティとバージョン互換性に注意する必要があります。

Pytorch torch 参考手册 Pytorch torch リファレンスマニュアル