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, resnet18

torch.hub.load('pytorch/vision', 'resnet50') 的执行流程:

  1. clone 仓库到本地缓存(已存在则直接复用)
  2. 在隔离命名空间里执行 hubconf.py
  3. 按名找到入口函数,剩余参数原样传入,返回其调用结果

入口函数内部通常构造 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.hubHuggingFace Hub
分发单位git 仓库,入口是代码模型 repo:config + 权重文件
权重格式入口代码自定,多为 .pthsafetensors 为主
安全模型加载即执行代码,靠 trust_repo 把关config 驱动,默认不执行模型代码
版本管理git tag/branchcommit + 量化/精度变体体系
主流场景视觉小模型、论文官方实现NLP/大模型、多模态生态事实标准

分工现状:CV 的 backbone(ResNet、ViT、DINOv2、SAM)用 torch.hub 仍是第一路径;LLM 及 transformers 生态几乎全走 HF。两者在 PyTorch 架构总览 的生态版图里各占一层,并不冲突——torch.hub 产出标准 nn.Module,之后权重持久化见 模型权重保存,跨框架导出部署见 ONNX 模型导出。