PyTorch 缓存文件下载路径修改方法

了解PyTorch默认下载预训练模型的路径及其对C盘空间的影响后，用户可能希望更改该路径以节省C盘空间。本文将详细介绍如何通过临时和永久方案修改PyTorch模型下载路径，使其存储在其他磁盘上。

背景知识：

PyTorch在下载预训练模型时，默认会将其存储在Windows的C:\Users\<username>\.cache路径下。这种设置可能导致C盘空间被占用数G。为了解决这个问题，我们需要了解PyTorch是如何处理模型下载路径的。

当pretrained参数设置为True时，PyTorch会调用torch.utils中的load_state_dict_from_url函数，该函数最终会调用torch.utils.model_zoo.load_url函数。其中，model_dir参数决定了模型的存储路径，其默认值取决于环境变量TORCH_HOME。在Windows上，TORCH_HOME默认指向%USERPROFILE%\.cache，这与Linux中的~路径对应。

修改方法：

1. 临时方案：

   通过在运行PyTorch代码时临时设置环境变量TORCH_HOME，可以将模型下载路径更改到其他位置。以下是实现此目的的JavaScript代码：

   import os
   os.environ['TORCH_HOME'] = 'E:/Data/torch-cache'

   执行此代码后，PyTorch会将预训练模型下载到E:/Data/torch-cache目录。

2. 永久方案：

   为了长期解决模型下载路径的问题，可以通过修改系统环境变量来设置TORCH_HOME。具体步骤如下：

   • 右键点击“此电脑”，选择“属性”。
   • 点击“高级系统设置”。
   • 在“系统属性”窗口中，选择“环境变量”。
   • 在“系统变量”部分，点击“新建”按钮。
   • 在“变量名”字段中输入TORCH_HOME，在“变量值”字段中输入新的路径，例如E:/Data/torch-cache。
   • 点击“确定”保存更改。

   完成上述步骤后，PyTorch将始终将预训练模型下载到指定的路径。

实际效果：

以下是修改环境变量后的PyTorch模型下载日志示例：

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

可以看到，模型成功下载到了新的路径E:/Data/torch-cache。

参考资料：

• 如何更改PyTorch预训练模型的下载路径
• 在PyTorch中将预训练模型下载到指定路径

通过上述方法，用户可以灵活地管理PyTorch预训练模型的下载路径，从而有效利用磁盘空间。