Ubuntu上PyTorch的模型保存与加载

在Ubuntu上使用PyTorch进行模型保存和加载是一个常见的任务。以下是详细的步骤和示例代码，帮助你完成这一过程。

安装PyTorch

首先，确保你已经安装了PyTorch。你可以使用以下命令来安装：

pip install torch torchvision

保存模型

假设你有一个简单的神经网络模型，你可以使用torch.save()函数来保存整个模型。

import torch
import torch.nn as nn

# 定义一个简单的神经网络模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc(x)
        return x

# 创建模型实例
model = SimpleNet()

# 假设模型已经训练好
# 保存模型
torch.save(model, 'model.pth')

加载模型

加载模型时，你可以使用torch.load()函数，并将其加载到一个新的模型实例中。

# 创建一个新的模型实例
model = SimpleNet()

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

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

# 现在你可以使用加载的模型进行预测
# 例如：
# input_tensor = torch.randn(1, 1, 28, 28)
# output = model(input_tensor)

注意事项

1. 模型状态字典：torch.save()保存的是整个模型的状态字典，而torch.load()加载的是状态字典。你需要将加载的状态字典赋值给模型的state_dict()。
2. 评估模式：在加载模型后，通常需要将模型设置为评估模式（model.eval()），以确保在推理过程中不会应用dropout等训练时才使用的层。
3. 设备兼容性：如果你在不同的设备上保存和加载模型（例如，在GPU上保存，在CPU上加载），可能需要使用map_location参数来指定加载位置。

# 在CPU上加载模型
model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))

完整示例

以下是一个完整的示例，包括模型定义、训练、保存和加载：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc(x)
        return x

# 创建模型实例
model = SimpleNet()

# 假设我们有一些数据
# input_tensor = torch.randn(64, 1, 28, 28)
# target = torch.randint(0, 10, (64,))

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型（这里省略训练过程）
# for epoch in range(num_epochs):
#     optimizer.zero_grad()
#     outputs = model(input_tensor)
#     loss = criterion(outputs, target)
#     loss.backward()
#     optimizer.step()

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

# 加载模型
model = SimpleNet()
model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
model.eval()

# 使用加载的模型进行预测
# input_tensor = torch.randn(1, 1, 28, 28)
# output = model(input_tensor)

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