【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(weightsmodels.ResNet50_Weights.DEFAULT) print(Model loaded successfully) except Exception as e: print(fError: {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(weightsmodels.ResNet50_Weights.DEFAULT) except Exception as e: print(fError: {e}) # urllib.error.URLError: urlopen error [Errno -3] Temporary failure in name resolution # 或 ConnectionError错误场景三不知道模型下载到了哪里import torchvision.models as models import torch # 加载模型 model models.resnet50(weightsmodels.ResNet50_Weights.DEFAULT) # 模型下载到了哪里 # 默认路径: ~/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth # 但很多开发者不知道这个路径 import os default_path os.path.expanduser(~/.cache/torch/hub/checkpoints/) print(fDefault download path: {default_path}) if os.path.exists(default_path): print(fFiles: {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(fHub 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_dirNone, ...): 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_okTrue) 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(fHub directory: {torch.hub.get_dir()}) # 输出: /data/models/torch/hub # 加载模型将下载到指定路径 model models.resnet50(weightsmodels.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(fHub directory: {torch.hub.get_dir()}) # 输出: /data/models/torch import torchvision.models as models model models.resnet50(weightsmodels.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_okTrue) if os.path.exists(target_path): print(fModel already exists: {target_path}) return target_path print(fDownloading model from {url} to {target_path}...) urllib.request.urlretrieve(url, target_path) print(fDownload 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(weightsNone) # 不下载权重 state_dict torch.load(target_path, map_locationcpu) 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_okTrue) def load_resnet(self, model_nameresnet50, **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(fUnknown model: {model_name}) model_path os.path.join(self.model_dir, filename) if not os.path.exists(model_path): raise FileNotFoundError( fModel file not found: {model_path}\n fPlease download it first. ) # 创建模型结构不加载预训练权重 model_fn getattr(models, model_name) model model_fn(weightsNone, **kwargs) # 加载权重 state_dict torch.load(model_path, map_locationcpu) 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(fAvailable 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_okTrue) # 元数据文件 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, indent2) 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(fUnknown model: {model_name}. fAvailable: {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(fModel already cached: {filepath}) return filepath print(fDownloading {model_name} from {url}...) print(fTarget path: {filepath}) # 下载 try: urllib.request.urlretrieve(url, filepath) file_size os.path.getsize(filepath) / (1024 * 1024) print(fDownload 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(fDownload 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(fUnknown 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( fModel file not found: {filepath}\n fSet download_if_missingTrue to download automatically. ) # 创建模型结构 model_fn getattr(models, model_name) model model_fn(weightsNone, **model_kwargs) # 加载权重 state_dict torch.load(filepath, map_locationcpu) 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(fRemoved: {filepath}) del self.metadata[model_name] self._save_metadata() else: print(fModel 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_okTrue) print(fTORCH_HOME set to: {path}) staticmethod def set_hub_dir(path: str): 设置 hub 目录 import torch torch.hub.set_dir(path) print(fHub 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(fPyTorch version: {torch.__version__}) print(fCUDA 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_dirtarget_dir) print(fDownloading {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(fFailed: {e}) failed.append(name) print() print( * 50) print(fDownload 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_dircustom_path) print(f\nCache directory: {manager.cache_dir}) print(fAvailable models: {manager.list_available_models()}) # 3. 下载一个模型 print(\n--- Downloading ResNet18 ---) try: model_path manager.download_model(resnet18) print(fModel saved to: {model_path}) except Exception as e: print(fDownload failed (expected in offline env): {e}) # 4. 加载模型 print(\n--- Loading Model ---) try: model manager.load_model(resnet18, download_if_missingTrue) print(fModel loaded successfully!) print(fModel type: {type(model).__name__}) print(fParameters: {sum(p.numel() for p in model.parameters()):,}) except Exception as e: print(fLoad 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_missingFalse) print(Offline load successful!) except FileNotFoundError as e: print(fOffline 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(weightsmodels.ResNet50_Weights.DEFAULT) model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # torchvision 0.13 (旧 API已弃用) model models.resnet50(pretrainedTrue)3. 权重文件的命名不同版本的 torchvision 可能为同一模型使用不同的权重文件名。下载时确保使用与你的 torchvision 版本匹配的 URL。4. map_location 的重要性加载模型时使用map_locationcpu可以避免在没有 GPU 的环境中报错# 安全的加载方式 state_dict torch.load(filepath, map_locationcpu) 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 中提前下载模型运行时离线加载在多用户环境中使用共享路径定期验证模型文件完整性通过本文的方法你可以完全控制预训练模型的存储位置适应各种部署环境和需求。