Ubuntu PyTorch如何进行模型保存和加载 - 问答

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

安装PyTorch

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

pip install torch torchvision

保存模型

假设你已经训练好了一个模型，可以使用torch.save()函数来保存模型。以下是一个简单的示例：

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

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

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

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

加载模型

加载模型时，可以使用torch.load()函数。需要注意的是，加载模型时需要确保模型类已经定义好，否则会报错。以下是一个示例：

import torch
import torch.nn as nn

# 定义与保存时相同的模型类
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

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

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

# 如果模型是在GPU上训练的，需要将其移动到CPU上
model.to('cpu')

# 现在可以使用模型进行推理
input_data = torch.randn(1, 10)
output = model(input_data)
print(output)

保存和加载模型的状态字典

另一种常见的方法是保存和加载模型的状态字典（state_dict），而不是整个模型。这种方法更加灵活，因为你可以只保存和加载特定的部分。以下是一个示例：

保存状态字典

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

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

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

# 假设模型已经训练好
# 保存状态字典
torch.save(model.state_dict(), 'model_state_dict.pth')

加载状态字典

import torch
import torch.nn as nn

# 定义与保存时相同的模型类
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

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

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

# 加载状态字典
model.load_state_dict(torch.load('model_state_dict.pth'))

# 如果模型是在GPU上训练的，需要将其移动到CPU上
model.to('cpu')

# 现在可以使用模型进行推理
input_data = torch.randn(1, 10)
output = model(input_data)
print(output)

通过以上步骤，你可以在Ubuntu上使用PyTorch轻松地保存和加载模型。希望这些信息对你有所帮助！

0 赞

0 踩