Linux环境中PyTorch模型保存与加载

在Linux环境中，使用PyTorch保存和加载模型非常简单。以下是一些基本步骤：

保存模型

1. 定义模型：首先，你需要定义一个PyTorch模型。

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. 保存模型：使用torch.save()函数保存整个模型或模型的状态字典。

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

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

加载模型

1. 加载整个模型：

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

2. 加载模型的状态字典：

# 首先定义与保存时相同的模型结构
model = MyModel()

# 然后加载状态字典
model.load_state_dict(torch.load('model_state_dict.pth'))

注意事项

• 模型结构：加载模型的状态字典时，必须确保模型结构与保存时相同。如果模型结构不同，可能会导致加载失败或出现错误。
• 设备兼容性：如果你在GPU上训练模型，但在CPU上加载模型，需要将模型移动到CPU。可以使用model.to('cpu')方法。

model = torch.load('model.pth', map_location=torch.device('cpu'))

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

示例代码

以下是一个完整的示例，展示了如何定义、保存和加载模型：

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()

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

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

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

通过这些步骤，你可以在Linux环境中轻松地保存和加载PyTorch模型。