Ubuntu上PyTorch模型如何保存与加载

在Ubuntu上保存和加载PyTorch模型通常涉及以下步骤：

保存模型

1. 定义模型：首先，你需要定义一个PyTorch模型。
2. 训练模型：使用训练数据集对模型进行训练。
3. 保存模型：训练完成后，你可以使用torch.save()函数来保存整个模型或者模型的状态字典。

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)

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

# 实例化模型
model = SimpleModel()

# 假设模型已经训练完成
# 保存整个模型
torch.save(model, 'model.pth')

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

加载模型

1. 加载模型：使用torch.load()函数来加载模型或模型的状态字典。
2. 实例化模型：如果只保存了状态字典，你需要先实例化模型，然后将状态字典加载到模型中。

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

# 如果只保存了状态字典，需要先实例化模型
model = SimpleModel()
# 加载状态字典
model.load_state_dict(torch.load('model_state_dict.pth'))

# 确保模型在评估模式
model.eval()

注意事项

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

  # 在CPU上加载模型
  model = torch.load('model.pth', map_location=torch.device('cpu'))
  
  # 或者加载状态字典并在特定设备上加载
  model = SimpleModel().to('cuda')  # 假设你想在GPU上加载
  model.load_state_dict(torch.load('model_state_dict.pth', map_location=torch.device('cuda')))
  

• 版本兼容性：确保保存模型和加载模型的PyTorch版本兼容。不同版本的PyTorch可能会有不同的模型保存格式。

• 安全性：从不可信来源加载模型时要小心，因为恶意代码可能会伪装成模型文件。

通过遵循这些步骤，你可以在Ubuntu上轻松地保存和加载PyTorch模型。