Ubuntu PyTorch如何进行模型保存

在Ubuntu上使用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')

这会将模型的所有参数和结构保存到一个文件中。

加载模型

1. 加载模型：使用torch.load()函数来加载模型。

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

2. 使用模型：加载后的模型可以直接用于推理或继续训练。

input_data = torch.randn(1, 10)  # 示例输入数据
output = loaded_model(input_data)
print(output)

注意事项

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

# 保存模型时指定设备为CPU
torch.save(model.cpu(), 'model.pth')

# 加载模型时指定设备为CPU
loaded_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', map_location=torch.device('cpu'))

# 使用加载的模型进行推理
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

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