torch.hub 是 PyTorch 内置的预训练模型分发机制:一行
torch.hub.load()从 GitHub 仓库拉取模型代码并返回带权重的实例。它解决「模型怎么到你手里」的问题——视觉小模型场景至今常用,而 NLP 与大模型生态的主流已转向 HuggingFace Hub。
机制:git 仓库 + hubconf.py
torch.hub 没有中央注册中心,分发单位就是普通 git 仓库,唯一约定是仓库根目录放一个 hubconf.py:
# 仓库根的 hubconf.py:声明依赖,暴露入口函数
dependencies = ['torch']
from torchvision.models import resnet50, resnet18torch.hub.load('pytorch/vision', 'resnet50') 的执行流程:
- clone 仓库到本地缓存(已存在则直接复用)
- 在隔离命名空间里执行
hubconf.py - 按名找到入口函数,剩余参数原样传入,返回其调用结果
入口函数内部通常构造 torch.nn 模型并用 load_state_dict_from_url 下载权重,所以返回的是普通 nn.Module——可训练、可保存、可导出,与手写模型无异。
基本用法
import torch
model = torch.hub.load('pytorch/vision', 'resnet50', weights='DEFAULT')
model.eval()
# 复现实验应钉住版本:'pytorch/vision:v0.15.2'
x = torch.randn(1, 3, 224, 224)
with torch.inference_mode():
logits = model(x) # (1, 1000) ImageNet logits帮助函数:
torch.hub.list('pytorch/vision') # 列出仓库全部入口
torch.hub.help('pytorch/vision', 'resnet50') # 查看入口签名与文档
torch.hub.load_state_dict_from_url(url) # 只下权重,不执行模型代码
torch.hub.load('local/dir', 'entry', source='local') # 加载本地仓库weights 参数演变
torchvision 自 0.13 起弃用 pretrained=True,改为 Weights 枚举,同一架构可挂多套权重:
from torchvision.models import resnet50, ResNet50_Weights
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2) # 显式选权重版本
# weights="DEFAULT" 取该模型当前最优枚举;pretrained=True 仍可跑,仅映射到默认权重并告警枚举同时携带 weights.transforms(),把 resize/normalize 与训练预处理对齐,避免手写出错。
缓存与刷新
| 需求 | 做法 |
|---|---|
| 查看缓存位置 | torch.hub.get_dir(),默认 ~/.cache/torch/hub |
| 重定向缓存 | 设 TORCH_HOME,缓存变为 $TORCH_HOME/hub |
| 强制更新 | force_reload=True,重新拉取仓库并执行 |
| 清理 | 直接删缓存目录,无注册表残留 |
仓库缓存在 <hub>/owner_repo_branch/,下载的权重在 <hub>/checkpoints/。默认命中缓存即不碰远端,远端更新不自动同步——复现实验务必钉住 tag 而非追 main。
安全:trust_repo
hubconf.py 在加载时会被执行,等于运行别人仓库里的任意代码。trust_repo 控制信任策略:"check" 记录仓库 commit hash,之后每次加载比对,不一致即报错;True 完全跳过校验。只对官方仓库与知名 lab 放行;对陌生 repo 设 trust_repo=True 等于自愿执行潜在恶意代码。
生态位:对比 HuggingFace Hub
| 维度 | torch.hub | HuggingFace Hub |
|---|---|---|
| 分发单位 | git 仓库,入口是代码 | 模型 repo:config + 权重文件 |
| 权重格式 | 入口代码自定,多为 .pth | safetensors 为主 |
| 安全模型 | 加载即执行代码,靠 trust_repo 把关 | config 驱动,默认不执行模型代码 |
| 版本管理 | git tag/branch | commit + 量化/精度变体体系 |
| 主流场景 | 视觉小模型、论文官方实现 | NLP/大模型、多模态生态事实标准 |
分工现状:CV 的 backbone(ResNet、ViT、DINOv2、SAM)用 torch.hub 仍是第一路径;LLM 及 transformers 生态几乎全走 HF。两者在 PyTorch 架构总览 的生态版图里各占一层,并不冲突——torch.hub 产出标准 nn.Module,之后权重持久化见 模型权重保存,跨框架导出部署见 ONNX 模型导出。