Pytorch 如何使用torch.hub.load加载本地模型

在本文中，我们将介绍如何使用Pytorch的torch.hub.load函数加载本地模型。torch.hub.load函数是Pytorch提供的一个便捷的方式，可以加载经过训练好的模型并在本地进行推理。

1. 了解torch.hub.load函数

torch.hub.load函数是Pytorch 1.1版本引入的一个重要特性。它允许我们通过指定模型的URL或者本地路径，快速加载模型进行后续的操作。通过该函数加载的模型，可以直接进行推理或者微调操作。

下面是torch.hub.load函数的基本用法：

torch.hub.load(repo_or_dir, model_name, *args, **kwargs)

其中，参数的含义如下：
– repo_or_dir：模型所在的仓库或者本地路径。
– model_name：模型的名称。
– args：模型初始化时的参数。
– kwargs：模型初始化时的关键字参数。

2. 加载本地模型

首先，我们需要下载并保存本地模型。PyTorch提供了许多预训练的模型可以供我们使用，这些模型通常具有出色的性能和泛化能力。例如，我们可以选择预训练的ResNet模型。

import torch
import torchvision.models as models

resnet18 = models.resnet18(pretrained=True)
torch.save(resnet18.state_dict(), 'resnet18.pth')

在上述代码中，我们使用torchvision库加载了一个预训练的ResNet-18模型，并将模型的参数保存到了resnet18.pth文件中。接下来，我们将使用torch.hub.load函数加载这个本地的ResNet模型。

import torch.hub

model = torch.hub.load('local/path/to/models', 'resnet18')

在上述代码中，我们将模型的本地路径传递给了torch.hub.load函数的第一个参数repo_or_dir，然后通过第二个参数model_name指定要加载的模型。这样，我们就成功地加载了本地的ResNet模型。

3. 加载带有初始化参数的本地模型

有时候，我们需要在加载模型时传递一些额外的初始化参数。这时，我们可以使用torch.hub.load函数的*args和**kwargs参数。

下面是一个加载带有额外初始化参数的本地模型的示例：

class MyModel(torch.nn.Module):
    def __init__(self, num_classes=1000, pretrained=False):
        super(MyModel, self).__init__()
        self.model = models.resnet18(pretrained=pretrained)
        self.fc = torch.nn.Linear(512, num_classes)

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

model = MyModel(pretrained=False)
torch.save(model.state_dict(), 'mymodel.pth')

model = torch.hub.load('local/path/to/models', 'mymodel', num_classes=10)