pytorch如何调用训练好的模型

1485
2024/4/11 13:29:04
栏目: 深度学习
开发者测试专用服务器限时活动,0元免费领,库存有限,领完即止! 点击查看>>

要调用训练好的模型,首先需要加载模型的参数,并将其应用到一个 PyTorch 模型中。以下是一个简单的示例代码,展示如何加载模型参数并将其应用到一个模型中:

import torch
import torch.nn as nn

# 定义一个简单的神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)  # 输入维度为 10,输出维度为 1

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

# 加载模型参数
model = SimpleModel()
model.load_state_dict(torch.load('model.pth'))

# 设置模型为 evaluation 模式
model.eval()

# 使用模型进行预测
input_data = torch.randn(1, 10)  # 生成一个随机输入数据
output = model(input_data)
print(output)

在这个示例中,首先定义了一个简单的神经网络模型 SimpleModel,然后加载了预训练好的模型参数,并将其应用到模型中。最后,使用模型进行预测并输出结果。

辰迅云「云服务器」,即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘,价格低至29元/月。点击查看>>

推荐阅读: PyTorch中怎么加速模型