Ubuntu上PyTorch的模型保存与加载方法 - 问答

在Ubuntu上使用PyTorch保存和加载模型主要涉及以下几个步骤：

保存模型

定义模型：首先，你需要定义你的模型。例如，一个简单的线性回归模型：

import torch
import torch.nn as nn

class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)

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

训练模型：在训练过程中，你可以使用标准的PyTorch训练循环。
保存模型：使用torch.save()函数来保存整个模型或模型的状态字典。

# 假设model是已经训练好的模型
torch.save(model, 'model.pth')  # 保存整个模型
# 或者
torch.save(model.state_dict(), 'model_state_dict.pth')  # 只保存模型的状态字典

加载模型

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

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

# 或者，如果你之前只保存了状态字典
model = LinearRegressionModel()  # 创建一个新的模型实例
model.load_state_dict(torch.load('model_state_dict.pth'))

注意：加载模型时，如果模型是在不同的环境中训练的（例如，使用了不同的PyTorch版本或不同的操作系统），可能会遇到兼容性问题。在这种情况下，你可能需要重新训练模型或使用map_location参数来指定加载模型的设备。

使用模型进行预测

加载模型后，你可以像平常一样使用它进行预测：

# 假设你有一个输入数据x
x = torch.tensor([[1.0]])
y_pred = model(x)
print(y_pred)

以上就是在Ubuntu上使用PyTorch保存和加载模型的基本方法。

0 赞

0 踩