修改 PyTorch Cache 文件下载路径

本文最后更新于:2022年8月25日 下午

PyTorch 自带很多预训练模型,在使用时会自动下载,本文记录修改下载位置的方法。

背景

  • PyTorch 下载预训练模型总得放个地方无可厚非,但默认路径在 Windows 中是 C:\Users\<username>\.cache ,很可能占用 C 盘几个 G 的空间,尝试修改该路径

模型加方式

  • pretrained为True时,PyTorch会调用torch.utilsload_state_dict_from_url函数
  • load_state_dict_from_url函数最终调用torch.utils.model_zoo.load_url函数
  • 其中的参数model_dir就是保存的目录,这里它默认会使用环境变量TORCH_HOME
  • 默认情况下环境变量TORCH_HOME的值为~/.cache
  • 在windows下就是%USERPROFILE%\.cache,其中%USERPROFILE%表示环境变量%USERPROFILE%的值,它对应Linux的~

修改方法

  • 搞清楚原理就容易修改了

临时方案

  • 运行 torch 时现场设置环境变量 ‘TORCH_HOME’
1
2
import os
os.environ['TORCH_HOME']='E:/Data/torch-cache'

永久方案

  • 修改环境变量,加入 ‘TORCH_HOME’

  • 之后下载模型就会到新路径去了
1
2
2022-08-25 14:53:36,183 - timm.models.helpers - INFO - Loading pretrained weights from url (https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/wide_resnet50_racm-8234f177.pth)
Downloading: "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/wide_resnet50_racm-8234f177.pth" to E:/Data/torch-cache\hub\checkpoints\wide_resnet50_racm-8234f177.pth

参考资料


修改 PyTorch Cache 文件下载路径
https://www.zywvvd.com/notes/study/deep-learning/pytorch/change-torch-default-cache-path/change-torch-default-cache-path/
作者
Yiwei Zhang
发布于
2022年8月25日
许可协议