轻松修改PyTorch预训练模型下载路径：一步一步指南

引言

在深度学习的研究和实践中，使用PyTorch框架及其预训练模型已成为一种常见做法。这些预训练模型不仅可以帮助我们快速启动项目，还能显著提升模型的初始性能。然而，PyTorch默认会将下载的预训练模型保存在用户的主目录下的特定文件夹中，这可能不是每个用户都想要或能接受的情况。例如，你可能希望将模型保存在一个专门的数据存储服务器或者具有更大磁盘空间的分区上。

幸运的是，PyTorch提供了灵活的方式来修改预训练模型的下载路径。下面，我将详细介绍如何做到这一点。

1. 理解PyTorch的模型加载机制

在PyTorch中，加载预训练模型通常是通过torchvision.models模块完成的，例如加载一个预训练的ResNet模型：

import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)

当pretrained=True时，PyTorch会尝试从Internet下载相应的预训练权重。下载的具体位置是由PyTorch内部管理的，但我们可以通过设置环境变量来更改它。

2. 修改环境变量以改变下载路径

PyTorch使用TORCH_HOME环境变量来指定默认的模型和数据下载路径。如果你没有设置这个环境变量，PyTorch会使用默认路径（通常是用户主目录下的.torch文件夹）。

要修改下载路径，你可以在Python脚本中或者在操作系统的环境变量设置中指定TORCH_HOME。

在Python脚本中设置

在Python脚本中，你可以使用os.environ来设置环境变量，但请注意，这仅对当前运行的Python进程有效。

import os
import torchvision.models as models
# 设置TORCH_HOME环境变量
os.environ['TORCH_HOME'] = '/path/to/your/custom/directory'
# 现在加载预训练模型
resnet18 = models.resnet18(pretrained=True)

在操作系统中设置

对于更持久的解决方案，你应该在操作系统的环境变量设置中添加或修改TORCH_HOME。这取决于你使用的操作系统。

• Windows：在“系统属性”的“高级”标签页中，点击“环境变量”按钮，然后添加或修改TORCH_HOME变量。

• Linux/macOS：在你的shell配置文件中（如.bashrc、.zshrc等），添加如下行：

  export TORCH_HOME=/path/to/your/custom/directory

  然后，重新加载配置文件或重新登录你的shell。

3. 验证更改

为了验证你是否成功更改了下载路径，你可以尝试加载一个预训练模型，并检查新路径下是否创建了相应的文件夹和文件。

4. 注意事项

• 修改TORCH_HOME会影响所有使用PyTorch下载的数据和模型，包括通过torchvision.datasets加载的数据集。
• 确保你的新路径是可写的，并且PyTorch有足够的权限访问它。
• 如果你在多个项目中需要不同的下载路径，考虑在Python脚本中动态设置TORCH_HOME，以避免全局更改。

结论

通过修改TORCH_HOME环境变量，你可以轻松地将PyTorch预训练模型的下载路径更改为你想要的任何位置。这不仅可以帮助你更好地管理你的数据和模型，还可以确保它们存储在适合你的工作流和存储需求的位置。希望这篇指南能帮助你实现这一目标！