python - 有什么方法可以将 PyTorch 中可用的预训练模型下载到特定路径？

Question

我指的是可以在这里找到的模型:https://pytorch.org/docs/stable/torchvision/models.html#torchvision-models

Answer 1

作为，@dennlinger在他的 answer 中提到: torch.utils.model_zoo , 在您加载预训练模型时被内部调用。

更具体地说，每次加载预训练模型时都会调用方法:torch.utils.model_zoo.load_url()。同样的文档提到:

The default value of model_dir is $TORCH_HOME/models where $TORCH_HOME defaults to ~/.torch.

The default directory can be overridden with the $TORCH_HOME environment variable.

这可以按如下方式完成:

import torch 
import torchvision
import os

# Suppose you are trying to load pre-trained resnet model in directory- models\resnet

os.environ['TORCH_HOME'] = 'models\\resnet' #setting the environment variable
resnet = torchvision.models.resnet18(pretrained=True)

我通过在 PyTorch 的 GitHub 存储库中提出问题来找到上述解决方案: https://github.com/pytorch/vision/issues/616

这导致了文档的改进，即上述解决方案。