这篇教程将PyTorch中可用的预训练模型下载到特定路径?写得很实用,希望能帮到您。
将PyTorch中可用的预训练模型下载到特定路径?
默认值model_dir 是$TORCH_HOME/models 其中 $TORCH_HOME 默认~/.torch 。
可以使用$TORCH_HOME 环境变量覆盖默认目录。
可以按照以下步骤进行:
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)
或者在ubuntu:设定下载存储的目录为vgg目录下
import torch
import torchvision
import os
os.environ['TORCH_HOME'] = 'vgg' #setting the environment variable
vgg = torchvision.models.vgg19(pretrained=True)
ACE算法--图像增强技术及实现代码 pytorch实现从本地加载 .pth 格式模型 |