Linux下PyTorch模型保存与加载

在Linux下，使用PyTorch保存和加载模型的步骤如下：

1. 保存模型

首先，我们需要定义一个PyTorch模型。这里以一个简单的多层感知器（MLP）为例：

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

input_size = 784
hidden_size = 128
output_size = 10

model = MLP(input_size, hidden_size, output_size)

接下来，我们可以使用torch.save()函数将模型保存到文件中：

torch.save(model, 'model.pth')

这将在当前目录下创建一个名为model.pth的文件，其中包含模型的权重和结构信息。

1. 加载模型

要加载保存的模型，我们首先需要重新创建相同的模型结构。然后，使用torch.load()函数从文件中加载模型的权重和结构信息，并将其应用到新创建的模型上：

# 重新创建相同的模型结构
model_loaded = MLP(input_size, hidden_size, output_size)

# 从文件中加载模型的权重和结构信息
model_loaded.load_state_dict(torch.load('model.pth'))

# 将模型设置为评估模式
model_loaded.eval()

现在，model_loaded变量包含了与原始模型相同的权重和结构信息，可以用于进行预测等操作。

注意：在加载模型时，确保使用与保存模型时相同的PyTorch版本，否则可能会出现兼容性问题。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！