PyTorch在Linux上的模型保存与加载方法

在PyTorch中，模型的保存和加载通常使用torch.save()和torch.load()函数。以下是在Linux上保存和加载模型的基本步骤：

保存模型

1. 定义模型：首先，你需要定义一个模型类并实例化它。

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 定义模型层
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = MyModel()

2. 训练模型：在训练过程中，模型的权重会不断更新。

# 假设这里有一些训练代码
for epoch in range(10):
    # 训练步骤...
    pass

3. 保存模型：使用torch.save()函数保存整个模型或仅保存模型的状态字典。

# 保存整个模型
torch.save(model, 'model.pth')

# 或者只保存模型的状态字典
torch.save(model.state_dict(), 'model_state_dict.pth')

加载模型

1. 加载模型：使用torch.load()函数加载模型。如果你之前保存了整个模型，可以直接加载；如果只保存了状态字典，则需要先创建一个模型实例，然后加载状态字典。

# 加载整个模型
model = torch.load('model.pth')

# 或者加载模型的状态字典
model = MyModel()  # 创建模型实例
model.load_state_dict(torch.load('model_state_dict.pth'))

注意事项

• 设备兼容性：如果你在不同的设备（CPU/GPU）上保存和加载模型，可能会遇到设备不匹配的问题。可以使用map_location参数来指定加载模型时的设备。

# 在CPU上加载模型
model = torch.load('model.pth', map_location=torch.device('cpu'))

# 或者在GPU上加载模型（假设你有可用的GPU）
model = torch.load('model.pth', map_location=torch.device('cuda'))

• 版本兼容性：不同版本的PyTorch可能会有不同的模型保存格式。确保在加载模型时使用与保存模型时相同或兼容的PyTorch版本。

通过以上步骤，你可以在Linux上轻松地保存和加载PyTorch模型。