工作原理解析

1. PyTorch的缓存机制

PyTorch在下载预训练模型时，会按照以下优先级顺序寻找缓存目录：

python

# PyTorch内部的逻辑（简化版）
def get_cache_dir():
    # 1. 首先检查 TORCH_HOME 环境变量
    if 'TORCH_HOME' in os.environ:
        return os.environ['TORCH_HOME']
    
    # 2. 如果没有，检查 XDG_CACHE_HOME（Linux标准）
    if 'XDG_CACHE_HOME' in os.environ:
        return os.path.join(os.environ['XDG_CACHE_HOME'], 'torch')
    
    # 3. 最后使用默认路径
    return os.path.expanduser('~/.cache/torch')  # Windows: C:\Users\用户名\.cache\torch

2. 环境变量的作用

当你设置 TORCH_HOME=D:\pytorch_models 时：

python

import os
import torch

# PyTorch会读取这个环境变量
cache_dir = os.environ.get('TORCH_HOME', default_path)
print(f"模型将下载到: {cache_dir}")

# 实际下载时，PyTorch内部会调用类似这样的代码
model_path = os.path.join(cache_dir, 'hub', 'checkpoints', 'vgg16-397923af.pth')

3. 完整的调用链路

python

# 你的代码
model = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1)

# ↓ PyTorch内部处理
# 1. 检查本地是否已有模型文件
# 2. 读取环境变量确定缓存目录
cache_dir = os.environ.get('TORCH_HOME', '~/.cache/torch')

# 3. 构建完整路径
model_file = os.path.join(cache_dir, 'hub', 'checkpoints', 'vgg16-397923af.pth')

# 4. 如果文件不存在，从网络下载到这个路径
if not os.path.exists(model_file):
    download_model(url, model_file)

4. 验证这个过程

你可以用这个代码验证：

python

import os
import torch

print("PyTorch会使用的缓存目录:")
print("TORCH_HOME环境变量:", os.environ.get('TORCH_HOME', '未设置'))
print("torch.hub实际使用的目录:", torch.hub.get_dir())

# 这会显示模型文件的完整路径
print("VGG16模型会保存在:", torch.hub.get_dir() + "/checkpoints/vgg16-397923af.pth")

5. 为什么这样设计？

• 灵活性: 用户可以自由选择存储位置
• 统一管理: 所有PyTorch模型都存在一个地方
• 跨项目共享: 多个项目可以共用同一个模型文件
• 磁盘空间优化: 避免重复下载相同模型

总结: 环境变量就像是给PyTorch一个"地址"，告诉它："嘿，把模型文件放到这个地方！"，PyTorch就会乖乖地按照这个地址去存储和查找模型文件。

那这个环境变量是不是只能设置这个名字

不完全是！环境变量的名字分为两类：

1. PyTorch框架专用的环境变量（固定名字）

这些是PyTorch内部硬编码识别的，名字不能改变：

python

# PyTorch预定义的环境变量（名字固定）
TORCH_HOME          # PyTorch模型缓存目录
CUDA_VISIBLE_DEVICES # 指定使用哪些GPU
PYTORCH_CUDA_ALLOC_CONF # CUDA内存分配配置
TORCH_CUDNN_V8_API_DISABLED # 禁用cuDNN v8 API
OMP_NUM_THREADS     # OpenMP线程数
MKL_NUM_THREADS     # Intel MKL线程数

2. 你自己定义的环境变量（任意名字）

你可以创建任何名字的环境变量：

python

import os

# 设置自定义环境变量
os.environ['MY_MODEL_PATH'] = 'D:/my_models'
os.environ['PROJECT_ROOT'] = 'D:/python_study'
os.environ['DATA_DIR'] = 'E:/datasets'

# 在代码中使用
model_path = os.environ.get('MY_MODEL_PATH', 'default_path')
data_dir = os.environ.get('DATA_DIR', './data')