Pytorch 如何修改 Torch Hub 模型下载的路径

在本文中，我们将介绍如何修改 Pytorch Torch Hub 模型下载的路径。Torch Hub 提供了一种方便的方式来下载和使用预训练的深度学习模型。默认情况下，Torch Hub 将模型下载到用户主目录下的“.torch”文件夹中。然而，有时我们希望将模型下载到不同的路径，例如外部硬盘驱动器或特定的文件夹。下面是一些示例说明，展示了如何修改 Torch Hub 模型下载的路径。

阅读更多：Pytorch 教程

方法一：使用环境变量

Pytorch 提供了一个名为TORCH_HOME的环境变量，可以用于指定 Torch Hub 模型下载的路径。我们可以将该环境变量设置为我们期望的路径，然后 Pytorch 将在该路径下下载和缓存模型。以下是设置环境变量的示例代码：

import os
os.environ['TORCH_HOME'] = '/path/to/download/models'

在这个示例中，我们将TORCH_HOME环境变量设置为/path/to/download/models。接下来，当我们使用 Torch Hub 下载模型时，Pytorch 将自动将模型保存到指定的路径中。

方法二：修改默认路径

如果我们希望在整个项目中使用相同的模型下载路径，我们可以修改 Pytorch 的默认路径。默认情况下，Pytorch 将模型保存到用户主目录下的“.torch”文件夹中。我们可以通过以下步骤修改默认路径：

首先，找到并打开.torch文件夹下的hubconf.py文件。该文件定义了 Torch Hub 的默认路径。

接下来，在hubconf.py文件中，找到以下代码：

torch_home = os.path.expanduser(
    os.getenv('TORCH_HOME', '~/.torch'))

将上述代码修改为我们期望的路径，例如：

torch_home = '/path/to/download/models'

保存文件后，Pytorch 的 Torch Hub 将自动将模型保存到指定的路径中。

需要注意的是，修改默认路径可能会影响其他使用 Torch Hub 的代码。确保在修改路径之前备份原始的hubconf.py文件。

方法三：手动指定路径

如果我们希望仅在特定的场景中更改模型下载路径，而不影响整个项目，我们可以在具体代码段中手动指定路径。以下是一个示例代码：

import torch
import torch.hub

model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet18', pretrained=True, force_reload=False, verbose=False, model_dir='/path/to/download/models')

在这个示例中，我们使用了 Torch Hub 下载了一个预训练的 ResNet-18 模型，并将模型保存到/path/to/download/models路径中。

通过手动指定路径，我们可以在特定的代码段中修改模型下载路径，而不影响其他地方的代码。