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

Question

我需要避免从网上下载模型（由于安装机器的限制）。

这可行，但从网上下载模型

model = torch.hub.load('pytorch/vision:v0.9.0', 'deeplabv3_resnet101', pretrained=True)

我已将.pth 文件和hubconf.py 文件放在/tmp/ 文件夹中，并将我的代码更改为

model = torch.hub.load('/tmp/', 'deeplabv3_resnet101', pretrained=True, source='local')

但令我惊讶的是，它仍然从互联网上下载模型。我究竟做错了什么？如何在本地加载模型？

为了给您提供更多细节，我在运行时具有只读卷的 Docker 容器中执行所有这些操作，这就是下载新文件失败的原因。

Answer 1

您可以采用两种方法在没有 Internet 连接的情况下在机器上获取可交付模型。

1. 在普通机器上加载带有预训练模型的 DeepLab，使用JIT 编译器将其导出为图形，然后放入机器中。脚本很容易理解：

    # To export
    model = torch.hub.load('pytorch/vision:v0.9.0', 'deeplabv3_resnet101', pretrained=True).eval()
    traced_graph = torch.jit.trace(model, torch.randn(1, 3, H, W))
    traced_graph.save('DeepLab.pth')
   
    # To load
    model = torch.jit.load('DeepLab.pth').eval().to(device)
   

   在这种情况下，权重和网络结构被保存为计算图，因此您不需要任何额外的文件。

2. 看看torchvision's GitHub repository。

   对于具有 Resnet101 主干权重的 DeepLabV3，有一个 download URL。

   您可以下载这些权重一次，然后使用带有 pretrained=False 标志的 torchvision 的 deeplab 并手动加载权重。

    model = torch.hub.load('pytorch/vision:v0.9.0', 'deeplabv3_resnet101', pretrained=False)
    model.load_state_dict(torch.load('downloaded weights path'))
   

   请注意，在 state dict 中可能有一个 ['state_dict'] 或一些类似的父键，您可以在其中使用：

    model.load_state_dict(torch.load('downloaded weights path')['state_dict'])