pytorch实现从本地加载 .pth 格式模型

1. 概述

在深度学习中，模型的训练通常需要付出较大的时间和计算资源。因此，为了方便使用训练好的模型，我们经常需要将模型保存在本地，以便在后续的预测或其他任务中使用。PyTorch是一种基于Python的用于构建深度学习模型的开源框架，本文将介绍如何从本地加载.pth格式的模型。

2. 加载.pth格式模型

2.1. 导入所需库

首先，我们需要导入PyTorch库，以及其它可能需要用到的库，如下所示：

import torch
import torchvision

2.2. 定义模型类

在加载.pth格式模型之前，我们需要先定义一个与保存的模型结构相同的模型类。这是因为在加载模型时，PyTorch需要知道如何构造模型的结构。如果保存时使用的是自定义的模型类，那么我们需要按照相同的方式定义模型类，并且确保保存和加载的模型类定义一致。

例如，假设我们的模型类定义如下：

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = torch.nn.Linear(3, 2)
        
    def forward(self, x):
        return self.fc(x)

2.3. 加载模型参数

模型的参数保存在.pth文件中，我们可以使用PyTorch提供的`load_state_dict`方法加载模型参数。具体步骤如下：

1. 创建一个模型对象：

model = MyModel()

2. 使用`torch.load`方法加载模型参数：

model.load_state_dict(torch.load('model.pth'))

3. 设置模型为评估模式：

model.eval()

通过以上步骤，我们就成功地从本地加载了.pth格式的模型。

3. 使用加载的模型进行预测

加载模型后，我们可以使用它来进行预测。下面是一个使用加载模型进行预测的示例：

# 定义输入数据
input_data = torch.tensor([[1, 2, 3]], dtype=torch.float32)
# 使用加载的模型进行预测
output = model(input_data)
# 打印预测结果
print(output)

在上述示例中，我们首先定义了一个输入张量`input_data`，然后使用加载的模型对输入进行预测，最后打印预测结果。

4. 总结

本文介绍了如何使用PyTorch从本地加载.pth格式的模型。首先，我们需要定义一个与保存的模型结构相同的模型类。然后，使用`load_state_dict`方法加载模型参数，并将模型设置为评估模式。最后，我们可以使用加载的模型进行预测。加载模型可以帮助我们节省模型训练的时间和计算资源，以便在后续的任务中快速使用已训练好的模型。