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

在Ubuntu上使用PyTorch进行模型的保存与加载，可以按照以下步骤操作：

模型保存

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

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)

2. 实例化模型并训练： 接下来，实例化模型并进行训练。这里我们假设你已经完成了训练过程。
3. 保存模型： 使用torch.save()函数来保存整个模型或模型的状态字典。

model = MyModel()
# 假设模型已经训练完成
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'))

注意：在加载模型时，确保模型类（在本例中为MyModel）已经在当前环境中定义。

2. 使用模型： 加载模型后，你可以像平常一样使用它进行推理或继续训练。

# 假设我们有一些输入数据
input_data = torch.randn(1, 10)

# 使用模型进行推理
output = model(input_data)
print(output)

注意事项

• 在保存和加载模型时，确保PyTorch的版本一致，以避免潜在的兼容性问题。
• 如果你在不同的环境中保存和加载模型（例如，在Ubuntu上训练，在Windows上加载），可能需要考虑模型文件路径的差异以及可能的操作系统兼容性问题。
• 为了确保模型的可移植性，建议将模型保存为状态字典的形式，并在加载时重新实例化模型类。