【Bug已解决】Is there any way I can download the pre-trained models available in PyTorch to a specific path? 解决方案
问题描述
在使用 PyTorch 的torchvision.models或torch.hub加载预训练模型时，模型文件默认下载到~/.cache/torch/hub/checkpoints/（或~/.torch/models/）目录。许多开发者需要将预训练模型下载到指定路径，原因包括：
- 服务器上用户主目录空间有限，需要将大模型存储到数据盘
- 多用户共享环境，需要将模型放在共享目录
- 离线环境，需要提前下载模型到指定位置
- CI/CD 环境中需要缓存模型以加速构建
- 自定义模型管理流程，需要统一存储路径
然而，PyTorch 并没有提供一个简单直接的"下载到指定路径"的参数。开发者经常遇到以下问题：
- 不知道如何更改默认下载路径
- 在离线环境中无法加载模型
- 多次下载同一模型浪费时间
- 在 Docker 容器中路径映射问题
TORCH_HOME环境变量不生效
本文将系统介绍多种将预训练模型下载到指定路径的方法，并提供完整的离线加载方案。
错误复现
错误场景一：默认下载路径不可写
import torchvision.models as models # 尝试加载预训练模型 try: model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) print("Model loaded successfully") except Exception as e: print(f"Error: {e}") # 可能的错误: # OSError: [Errno 30] Read-only file system: '/root/.cache/torch/hub/checkpoints' # 或 PermissionError: [Errno 13] Permission denied错误场景二：离线环境无法下载
import torchvision.models as models # 在没有网络的环境中 try: model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) except Exception as e: print(f"Error: {e}") # urllib.error.URLError: <urlopen error [Errno -3] Temporary failure in name resolution> # 或 ConnectionError错误场景三：不知道模型下载到了哪里
import torchvision.models as models import torch # 加载模型 model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) # 模型下载到了哪里？ # 默认路径: ~/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth # 但很多开发者不知道这个路径 import os default_path = os.path.expanduser('~/.cache/torch/hub/checkpoints/') print(f"Default download path: {default_path}") if os.path.exists(default_path): print(f"Files: {os.listdir(default_path)}")错误场景四：TORCH_HOME 设置不生效
import os # 设置环境变量（但可能在 import torch 之后设置，不生效） os.environ['TORCH_HOME'] = '/data/models/torch' import torch import torchvision.models as models # 检查实际路径 print(f"Hub dir: {torch.hub.get_dir()}") # 可能仍然是 ~/.cache/torch/hub 而不是 /data/models/torch/hub根因分析
1. PyTorch 的模型下载机制
PyTorch 加载预训练模型的过程：
- 检查本地缓存：首先检查
TORCH_HOME/hub/checkpoints/目录下是否已有模型文件 - 下载模型：如果本地没有，从远程 URL 下载
- 加载模型：将下载的权重文件加载到模型结构中
默认的TORCH_HOME路径：
- Linux/macOS:
~/.cache/torch/ - Windows:
C:\Users\<username>\.cache\torch\
2. 环境变量的读取时机
TORCH_HOME环境变量在import torch时被读取。如果在import torch之后设置环境变量，不会生效。
# 错误：先 import 再设置 import torch os.environ['TORCH_HOME'] = '/data/models' # 太晚了！ # 正确：先设置再 import import os os.environ['TORCH_HOME'] = '/data/models' import torch3. torchvision.models 的权重下载
torchvision.models中的预训练模型权重存储在 GitHub Releases 或其他 CDN 上。每个模型有对应的 URL，下载后缓存到checkpoints子目录。
# torchvision 内部的下载逻辑（简化版） def load_state_dict_from_url(url, model_dir=None, ...): if model_dir is None: # 使用默认路径: TORCH_HOME/hub/checkpoints hub_dir = torch.hub.get_dir() model_dir = os.path.join(hub_dir, 'checkpoints') os.makedirs(model_dir, exist_ok=True) filename = os.path.basename(url) cached_file = os.path.join(model_dir, filename) if not os.path.exists(cached_file): # 下载文件 download_url_to_file(url, cached_file, ...) return torch.load(cached_file)4. torch.hub 的机制
torch.hub.load()用于加载 GitHub 上的模型仓库，下载的代码缓存到TORCH_HOME/hub/目录，权重文件缓存到checkpoints子目录。
解决方案
方案一：设置 TORCH_HOME 环境变量
import os # 必须在 import torch 之前设置 os.environ['TORCH_HOME'] = '/data/models/torch' import torch import torchvision.models as models # 验证路径 print(f"Hub directory: {torch.hub.get_dir()}") # 输出: /data/models/torch/hub # 加载模型（将下载到指定路径） model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) print("Model loaded!")也可以在命令行中设置：
export TORCH_HOME=/data/models/torch python your_script.py方案二：使用 torch.hub.set_dir()
import torch # 在代码中设置 hub 目录 custom_dir = '/data/models/torch' torch.hub.set_dir(custom_dir) print(f"Hub directory: {torch.hub.get_dir()}") # 输出: /data/models/torch import torchvision.models as models model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)方案三：手动下载并加载
import torch import torchvision.models as models import os import urllib.request def download_model_to_path(url, target_path): """手动下载模型到指定路径""" os.makedirs(os.path.dirname(target_path), exist_ok=True) if os.path.exists(target_path): print(f"Model already exists: {target_path}") return target_path print(f"Downloading model from {url} to {target_path}...") urllib.request.urlretrieve(url, target_path) print(f"Download complete: {target_path}") return target_path # ResNet50 的权重 URL resnet50_url = "https://download.pytorch.org/models/resnet50-0676ba61.pth" target_path = "/data/models/pretrained/resnet50.pth" # 下载 download_model_to_path(resnet50_url, target_path) # 手动加载 model = models.resnet50(weights=None) # 不下载权重 state_dict = torch.load(target_path, map_location='cpu') model.load_state_dict(state_dict) print("Model loaded from custom path!")方案四：离线加载方案
import torch import torchvision.models as models import os class OfflineModelLoader: """ 离线模型加载器。 从指定目录加载预训练模型，不依赖网络。 """ def __init__(self, model_dir='/data/models/pretrained'): self.model_dir = model_dir os.makedirs(model_dir, exist_ok=True) def load_resnet(self, model_name='resnet50', **kwargs): """加载 ResNet 模型""" # 模型名到文件名的映射 filename_map = { 'resnet18': 'resnet18-f37072fd.pth', 'resnet34': 'resnet34-b627a593.pth', 'resnet50': 'resnet50-0676ba61.pth', 'resnet101': 'resnet101-63fe2227.pth', 'resnet152': 'resnet152-394f9c45.pth', } filename = filename_map.get(model_name) if not filename: raise ValueError(f"Unknown model: {model_name}") model_path = os.path.join(self.model_dir, filename) if not os.path.exists(model_path): raise FileNotFoundError( f"Model file not found: {model_path}\n" f"Please download it first." ) # 创建模型结构（不加载预训练权重） model_fn = getattr(models, model_name) model = model_fn(weights=None, **kwargs) # 加载权重 state_dict = torch.load(model_path, map_location='cpu') model.load_state_dict(state_dict) return model def list_available_models(self): """列出本地可用的模型""" if os.path.exists(self.model_dir): return os.listdir(self.model_dir) return [] # 使用 loader = OfflineModelLoader('/data/models/pretrained') print(f"Available models: {loader.list_available_models()}") model = loader.load_resnet('resnet50')完整修复代码
以下是一个完整的预训练模型管理器，支持下载到指定路径、离线加载、缓存管理：
import torch import torchvision.models as models import os import urllib.request import hashlib from typing import Optional, Dict, List import json # ============================================================ # 1. 预训练模型管理器 # ============================================================ class PretrainedModelManager: """ 预训练模型管理器。 - 支持下载到指定路径 - 支持离线加载 - 支持缓存管理 - 支持多模型管理 """ # torchvision 预训练模型的 URL 映射 MODEL_URLS = { 'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth', 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 'mobilenet_v2': 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth', 'mobilenet_v3_large': 'https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth', 'efficientnet_b0': 'https://download.pytorch.org/models/efficientnet_b0_rwightman-3dd342df.pth', 'efficientnet_b7': 'https://download.pytorch.org/models/efficientnet_b7_lukemelas-dcc49843.pth', 'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth', 'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth', } def __init__(self, cache_dir: str = None): """ Args: cache_dir: 模型缓存目录。默认为 ~/.cache/torch/hub/checkpoints """ if cache_dir is None: cache_dir = os.path.join(os.path.expanduser('~'), '.cache', 'torch', 'hub', 'checkpoints') self.cache_dir = cache_dir os.makedirs(self.cache_dir, exist_ok=True) # 元数据文件 self.metadata_file = os.path.join(self.cache_dir, 'model_metadata.json') self.metadata = self._load_metadata() def _load_metadata(self) -> Dict: """加载元数据""" if os.path.exists(self.metadata_file): with open(self.metadata_file, 'r') as f: return json.load(f) return {} def _save_metadata(self): """保存元数据""" with open(self.metadata_file, 'w') as f: json.dump(self.metadata, f, indent=2) def download_model(self, model_name: str, force: bool = False) -> str: """ 下载预训练模型到缓存目录。 Args: model_name: 模型名称 force: 是否强制重新下载 Returns: 模型文件路径 """ if model_name not in self.MODEL_URLS: raise ValueError(f"Unknown model: {model_name}. " f"Available: {list(self.MODEL_URLS.keys())}") url = self.MODEL_URLS[model_name] filename = os.path.basename(url) filepath = os.path.join(self.cache_dir, filename) if os.path.exists(filepath) and not force: print(f"Model already cached: {filepath}") return filepath print(f"Downloading {model_name} from {url}...") print(f"Target path: {filepath}") # 下载 try: urllib.request.urlretrieve(url, filepath) file_size = os.path.getsize(filepath) / (1024 * 1024) print(f"Download complete! Size: {file_size:.1f} MB") # 计算文件哈希 file_hash = self._compute_hash(filepath) # 更新元数据 self.metadata[model_name] = { 'url': url, 'filename': filename, 'filepath': filepath, 'size_mb': round(file_size, 1), 'hash': file_hash, 'download_time': os.path.getmtime(filepath), } self._save_metadata() return filepath except Exception as e: print(f"Download failed: {e}") if os.path.exists(filepath): os.remove(filepath) raise def _compute_hash(self, filepath: str, algorithm: str = 'sha256') -> str: """计算文件哈希""" hasher = hashlib.new(algorithm) with open(filepath, 'rb') as f: while True: chunk = f.read(65536) if not chunk: break hasher.update(chunk) return hasher.hexdigest() def load_model(self, model_name: str, download_if_missing: bool = True, **model_kwargs) -> torch.nn.Module: """ 加载预训练模型。 Args: model_name: 模型名称 download_if_missing: 如果本地没有是否自动下载 model_kwargs: 模型参数 Returns: 加载好权重的模型 """ if model_name not in self.MODEL_URLS: raise ValueError(f"Unknown model: {model_name}") url = self.MODEL_URLS[model_name] filename = os.path.basename(url) filepath = os.path.join(self.cache_dir, filename) if not os.path.exists(filepath): if download_if_missing: self.download_model(model_name) else: raise FileNotFoundError( f"Model file not found: {filepath}\n" f"Set download_if_missing=True to download automatically." ) # 创建模型结构 model_fn = getattr(models, model_name) model = model_fn(weights=None, **model_kwargs) # 加载权重 state_dict = torch.load(filepath, map_location='cpu') model.load_state_dict(state_dict) return model def list_cached_models(self) -> List[Dict]: """列出已缓存的模型""" cached = [] for model_name, info in self.metadata.items(): filepath = info.get('filepath', '') exists = os.path.exists(filepath) cached.append({ 'name': model_name, 'filename': info.get('filename', ''), 'size_mb': info.get('size_mb', 0), 'exists': exists, 'hash': info.get('hash', '')[:16] + '...', }) return cached def list_available_models(self) -> List[str]: """列出所有可用的模型""" return list(self.MODEL_URLS.keys()) def remove_model(self, model_name: str): """删除缓存的模型""" if model_name in self.metadata: filepath = self.metadata[model_name].get('filepath') if filepath and os.path.exists(filepath): os.remove(filepath) print(f"Removed: {filepath}") del self.metadata[model_name] self._save_metadata() else: print(f"Model not in metadata: {model_name}") def clear_cache(self): """清空所有缓存""" for model_name in list(self.metadata.keys()): self.remove_model(model_name) print("Cache cleared.") def get_cache_size(self) -> float: """获取缓存总大小（MB）""" total = 0 for info in self.metadata.values(): filepath = info.get('filepath', '') if os.path.exists(filepath): total += os.path.getsize(filepath) return total / (1024 * 1024) def verify_model(self, model_name: str) -> bool: """验证模型文件完整性""" if model_name not in self.metadata: return False info = self.metadata[model_name] filepath = info.get('filepath', '') if not os.path.exists(filepath): return False current_hash = self._compute_hash(filepath) stored_hash = info.get('hash', '') return current_hash == stored_hash # ============================================================ # 2. 环境配置工具 # ============================================================ class TorchEnvConfig: """ PyTorch 环境配置工具。 """ @staticmethod def set_torch_home(path: str): """设置 TORCH_HOME 环境变量""" os.environ['TORCH_HOME'] = path os.makedirs(os.path.join(path, 'hub', 'checkpoints'), exist_ok=True) print(f"TORCH_HOME set to: {path}") @staticmethod def set_hub_dir(path: str): """设置 hub 目录""" import torch torch.hub.set_dir(path) print(f"Hub directory set to: {torch.hub.get_dir()}") @staticmethod def get_current_paths(): """获取当前路径配置""" import torch return { 'TORCH_HOME': os.environ.get('TORCH_HOME', 'Not set'), 'hub_dir': torch.hub.get_dir(), 'checkpoints_dir': os.path.join(torch.hub.get_dir(), 'checkpoints'), } @staticmethod def print_info(): """打印当前环境信息""" import torch paths = TorchEnvConfig.get_current_paths() print("=" * 50) print("PyTorch Environment Info") print("=" * 50) print(f"PyTorch version: {torch.__version__}") print(f"CUDA available: {torch.cuda.is_available()}") for key, value in paths.items(): print(f"{key}: {value}") print("=" * 50) # ============================================================ # 3. 批量下载工具 # ============================================================ def batch_download_models(model_names: List[str], target_dir: str): """ 批量下载多个预训练模型。 Args: model_names: 模型名称列表 target_dir: 目标目录 """ manager = PretrainedModelManager(cache_dir=target_dir) print(f"Downloading {len(model_names)} models to: {target_dir}\n") success = [] failed = [] for name in model_names: try: print(f"--- {name} ---") manager.download_model(name) success.append(name) except Exception as e: print(f"Failed: {e}") failed.append(name) print() print("=" * 50) print(f"Download Summary:") print(f" Success: {len(success)} - {success}") print(f" Failed: {len(failed)} - {failed}") print(f" Total cache size: {manager.get_cache_size():.1f} MB") return success, failed # ============================================================ # 4. 主程序 # ============================================================ def main(): print("PyTorch Pretrained Model Path Management") print("=" * 60) # 1. 显示当前环境 TorchEnvConfig.print_info() # 2. 创建模型管理器（指定自定义路径） custom_path = os.path.expanduser('~/my_models/torch_pretrained') manager = PretrainedModelManager(cache_dir=custom_path) print(f"\nCache directory: {manager.cache_dir}") print(f"Available models: {manager.list_available_models()}") # 3. 下载一个模型 print("\n--- Downloading ResNet18 ---") try: model_path = manager.download_model('resnet18') print(f"Model saved to: {model_path}") except Exception as e: print(f"Download failed (expected in offline env): {e}") # 4. 加载模型 print("\n--- Loading Model ---") try: model = manager.load_model('resnet18', download_if_missing=True) print(f"Model loaded successfully!") print(f"Model type: {type(model).__name__}") print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}") except Exception as e: print(f"Load failed: {e}") # 5. 列出缓存 print("\n--- Cached Models ---") cached = manager.list_cached_models() for m in cached: print(f" {m['name']}: {m['size_mb']}MB, exists={m['exists']}") print(f"\nTotal cache size: {manager.get_cache_size():.1f} MB") # 6. 验证模型 print("\n--- Model Verification ---") for name in [m['name'] for m in cached]: valid = manager.verify_model(name) print(f" {name}: {'Valid' if valid else 'Invalid'}") # 7. 演示离线加载 print("\n--- Offline Loading Demo ---") try: model = manager.load_model('resnet18', download_if_missing=False) print("Offline load successful!") except FileNotFoundError as e: print(f"Offline load failed: {e}") if __name__ == '__main__': main()常见陷阱与注意事项
1. 环境变量设置时机
TORCH_HOME必须在import torch之前设置。如果在 Jupyter Notebook 中，需要在第一个单元格中设置，然后重启内核。
# 正确 import os os.environ['TORCH_HOME'] = '/data/models' import torch # 错误 import torch os.environ['TORCH_HOME'] = '/data/models' # 不生效2. torchvision 版本差异
不同版本的 torchvision 使用不同的 API 加载预训练权重：
# torchvision >= 0.13 (新 API) model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) # torchvision < 0.13 (旧 API，已弃用) model = models.resnet50(pretrained=True)3. 权重文件的命名
不同版本的 torchvision 可能为同一模型使用不同的权重文件名。下载时确保使用与你的 torchvision 版本匹配的 URL。
4. map_location 的重要性
加载模型时使用map_location='cpu'可以避免在没有 GPU 的环境中报错：
# 安全的加载方式 state_dict = torch.load(filepath, map_location='cpu') model.load_state_dict(state_dict)5. Docker 中的路径映射
在 Docker 中使用自定义模型路径时，需要将宿主机目录映射到容器中：
docker run -v /host/models:/container/models \ -e TORCH_HOME=/container/models/torch \ your_image6. 多进程环境下的下载竞争
在多进程训练中，多个进程可能同时尝试下载同一模型。使用文件锁或提前下载来避免竞争：
import fcntl def safe_download(url, filepath): lock_file = filepath + '.lock' with open(lock_file, 'w') as f: fcntl.flock(f, fcntl.LOCK_EX) # 排他锁 if not os.path.exists(filepath): urllib.request.urlretrieve(url, filepath) fcntl.flock(f, fcntl.LOCK_UN)7. Hugging Face 模型的路径
Hugging Face Transformers 使用不同的缓存机制：
from transformers import AutoModel # 设置缓存目录 model = AutoModel.from_pretrained('bert-base-uncased', cache_dir='/data/models/hf')8. 模型哈希验证
下载后验证模型文件的哈希值，确保文件完整：
import hashlib def verify_file(filepath, expected_hash): hasher = hashlib.sha256() with open(filepath, 'rb') as f: for chunk in iter(lambda: f.read(65536), b''): hasher.update(chunk) return hasher.hexdigest() == expected_hash9. 模型版本管理
不同版本的预训练权重可能有不同的性能表现。建议记录使用的权重版本：
# 记录模型版本 model_info = { 'name': 'resnet50', 'weights': 'IMAGENET1K_V1', 'url': 'https://download.pytorch.org/models/resnet50-0676ba61.pth', 'hash': 'sha256:...', }总结
将 PyTorch 预训练模型下载到指定路径的方法：
设置
TORCH_HOME环境变量：最简单的方法，在import torch之前设置。影响所有 PyTorch 模型的下载路径。使用
torch.hub.set_dir()：在代码中动态设置，适合需要根据条件选择路径的场景。手动下载并加载：最灵活的方法，完全控制下载路径和加载过程。适合离线环境和自定义管理。
使用模型管理器：本文提供的
PretrainedModelManager类封装了下载、加载、缓存管理、完整性验证等功能，适合生产环境使用。
最佳实践：
- 在项目初始化时统一设置模型路径
- 使用元数据文件记录模型版本和哈希
- 在 CI/CD 中提前下载模型，运行时离线加载
- 在多用户环境中使用共享路径
- 定期验证模型文件完整性
通过本文的方法，你可以完全控制预训练模型的存储位置，适应各种部署环境和需求。
