PyTorch预训练模型下载与缓存路径管理:以VGG16为例 很多人第一次接触PyTorch预训练模型都是从一行models.vgg16(weights...)开始的。我当年也一样在服务器上敲完这行命令看着终端里那个下载进度条一路跑完就理所当然地认为权重文件落在了当前目录。后来要给项目做迁移找遍了工作目录也没看到那个.pth最后才发现默认缓存目录在~/.cache/torch/checkpoints而且那里已经默默躺了好几个GB的模型文件。这篇就以VGG为例子把预训练模型下载和加载的过程完整拆一遍默认缓存是怎么运作的、文件名里的哈希有什么用、自定义保存路径有哪几种打法、以及下载过程中常见的坑该怎么处理。不管是刚入门想在本地跑通VGG还是在服务器、离线环境、团队共享场景下需要管理一批预训练权重这篇应该都能帮你省下不少时间。1. 权重默认落在哪个目录几条规则先讲清楚1.1 首次调用触发下载之后走缓存先理解一个最基本的机制torchvision.models.vgg16(weightsVGG16_Weights.IMAGENET1K_V1)这行代码在你第一次执行时会去官方模型站下载对应的权重文件随后把文件存到本地缓存目录等第二次再执行同样的代码检测到目标文件已经存在就直接读缓存不再联网。默认的缓存根目录是~/.cache/torch也就是家目录下的隐藏路径。权重文件会放在它的checkpoints子目录里所以完整的默认路径是~/.cache/torch/checkpoints。如果你设置了环境变量TORCH_HOME那根目录就变成$TORCH_HOME权重依旧在$TORCH_HOME/checkpoints下面。很多老教程会提到~/.torch/models那是很早期版本的习惯路径现在的PyTorch已经改成上面这套了。如果你照着老文章去找文件找不到别怀疑代码写错先看一眼版本。1.2 文件名里的那串哈希不是乱码打开~/.cache/torch/checkpoints你会看到类似vgg16-397923af.pth这样的文件。注意中间那串397923af它不是随便加的版本号而是整个权重文件MD5值的前8位十六进制字符。这串哈希的作用有两个。第一做内容校验PyTorch下载完成后可以通过比对哈希来判断文件是否完整。第二做缓存区分不同来源、不同版本的权重即使模型结构同名文件也不会互相覆盖。所以你会发现vgg16-397923af.pth和带BatchNorm的vgg16_bn-6c64b313.pth文件名并不一样。这个细节在后面手动搬运文件时特别关键如果你把文件改了名字PyTorch就不认这个缓存下次照样重新下载一遍。1.3 为什么VGG尤其需要关心路径拿VGG16来说它的参数量大约是1.38亿权重文件体积在528MB左右。这不是一个小数目。在本地个人电脑上误存到~/.cache可能只是占用一点磁盘但在多人共用的服务器上家目录配额往往很紧要是谁都不管三七二十一让所有模型默认下载到各自家目录几台机器合起来就是几十GB的重复垃圾。而且VGG系列有VGG11、VGG13、VGG16、VGG19好几种规格加上BN版本一次性把变体都下一遍光VGG就得好几个GB。如果项目里还有ResNet、EfficientNet这些模型默认路径下的文件会越堆越多。所以从第一天开始就规划好权重存放目录是省心而不是折腾。1.4 torch.hub.set_dir管不到torchvision权重这是一个非常常见的误区。很多人搜自定义路径时看到torch.hub.set_dir()以为设置一次就能让models.vgg16(weights...)的权重也改到新目录结果发现文件还是下到了老地方。原因是torch.hub.set_dir()只影响torch.hub.load()这种按repo加载模型时使用的hub缓存目录而torchvision.models里的权重下载走的是另一条路径——它内部调用的torch.hub.load_state_dict_from_url默认目录是根据TORCH_HOME计算出来的checkpoints目录。所以想让torchvision的权重改道你真正应该动的是TORCH_HOME或者直接给load_state_dict_from_url传model_dir参数这个下面细说。2. 自定义路径的四种思路原理、代码与选型2.1 方案一TORCH_HOME环境变量一句话全局改道最省事的做法是在运行代码之前设置环境变量export TORCH_HOME/data/models/torch然后正常写代码就行import torch from torchvision import models model models.vgg16(weightsmodels.VGG16_Weights.IMAGENET1K_V1)执行之后权重会保存到/data/models/torch/checkpoints/vgg16-397923af.pth。这个方案的好处是业务代码一行都不用改所有走默认机制的下载包括torch.hub.load都会跟着新目录走。代价是它是进程级/会话级的全局设置如果你在脚本里想为不同任务指定不同目录就不方便了。如果你习惯直接在Python脚本里控制可以在文件最顶部写上import os os.environ[TORCH_HOME] /data/models/torch注意这里有一个顺序问题只要这行代码在真正触发下载的语句之前执行都来得及因为PyTorch是在下载那一刻才去读环境变量的。但为了保险我还是建议放到脚本最前面别和一堆import混在一起否则排查起来很费劲。2.2 方案二手动下载到任意目录加载时指定如果你希望完全掌控文件位置甚至想给权重文件改名、放到项目自己的目录结构里那就不要用自动下载改成手动下载加手动加载。先用Python打印出官方权重URLfrom torchvision.models import VGG16_Weights weights VGG16_Weights.IMAGENET1K_V1 print(weights.url) # https://download.pytorch.org/models/vgg16-397923af.pth然后用wget或者浏览器把这个文件下载到任意目录比如/data/weights/my_vgg16.pthwget https://download.pytorch.org/models/vgg16-397923af.pth -O /data/weights/my_vgg16.pth接着在代码里这样加载import torch from torchvision import models model models.vgg16(weightsNone) # 不加载预训练构建空模型 state_dict torch.load(/data/weights/my_vgg16.pth, map_locationcpu) model.load_state_dict(state_dict) # 自己把权重塞进去这里weightsNone的意思是只搭VGG16的网络结构不触发任何下载。torch.load把文件读成字典load_state_dict再逐层对应到模型上。map_locationcpu是好习惯避免在GPU机器上保存的权重换到CPU机器时出现设备不匹配的问题。这个方案的理解成本最低而且和PyTorch的缓存机制彻底解耦。缺点是你需要自己维护URL和版本加载时多做两步。适合离线部署、或者想统一管理一批手动收集的权重文件的场景。2.3 方案三load_state_dict_from_url精准指定model_dir方案一和方案二之间还有一个折中不用环境变量而是直接调用下载函数并且指定存放目录。import torch from torchvision import models from torchvision.models import VGG16_Weights url VGG16_Weights.IMAGENET1K_V1.url state_dict torch.hub.load_state_dict_from_url( url, model_dir/data/models/torch/checkpoints, progressTrue ) model models.vgg16(weightsNone) model.load_state_dict(state_dict)这里的关键是model_dir参数。torchvision.models内部其实也是用这个函数去下载权重的只不过没有把目录参数暴露给你。你现在绕到这一层去调它就能只对这一次调用指定目录既不污染全局环境变量也能控制下载进度条、是否校验哈希等行为。需要注意文件名的规则仍然有效它会把URL末尾的vgg16-397923af.pth作为文件名保存到model_dir。如果你把model_dir设置成/data/models/torch/checkpoints那就正好和方案一中TORCH_HOME/data/models/torch的效果一致两者可以混用。2.4 方案四预置缓存目录提前把下载做掉这个方法本质上不是下载时怎么改路径而是让已经下载好的权重正好出现在PyTorch会去找的位置。适用于离线机器、内网环境以及团队里想让所有人共用同一份权重的情况。做法很简单在一台有网的机器上用默认方式下载一次然后把checkpoints目录里的文件原封不动拷到目标机器的对应位置。目标机器上只要设置了TORCH_HOME并且文件路径是$TORCH_HOME/checkpoints/vgg16-397923af.pth运行时就不会联网直接读这个文件。要确认目标路径可以用这个Python命令打印import os import torch torch_home os.environ.get(TORCH_HOME, os.path.expanduser(~/.cache/torch)) print(os.path.join(torch_home, checkpoints))这个方案看起来很朴素但在实际工程里是最高效的下载只发生一次之后的加载全部命中缓存速度快、不依赖外网也不担心某次下载中途失败。2.5 四种方案放在一起怎么选方案改动形式首次需要联网适合场景TORCH_HOME环境变量全局/进程级是单机、服务器统一缓存目录手动下载load_state_dict每次加载逻辑下载时需要离线部署、自定义命名、精细控制load_state_dict_from_url指定model_dir单次调用是不动全局又想改目录预置缓存目录全局/部署时仅准备阶段需要团队共享、内网、离线集群我的建议是个人开发用方案一最省事部署到离线环境用方案二或方案四需要在一段代码里对不同模型用不同目录时用方案三临时穿针引线。3. VGG16完整实操从配置路径到跑通推理3.1 环境与目录准备先确认环境里的torch和torchvision版本。weights这个新API从torchvision 0.13左右开始稳定提供如果你还在用特别老的版本建议先升级。验证命令python -c import torch, torchvision; print(torch.__version__, torchvision.__version__)然后创建目标目录并设置环境变量mkdir -p /data/models/torch/checkpoints export TORCH_HOME/data/models/torch如果你想每次登录都自动生效就把export TORCH_HOME/data/models/torch写进~/.bashrc或~/.profile。如果在容器里用可以在镜像的启动脚本里加上这一行。团队协作时我建议把环境变量的设置放到每个任务脚本的前三行而不是依赖个人shell配置这样换谁跑结果都一样。3.2 下载加载与状态检查写一个最简脚本import os os.environ[TORCH_HOME] /data/models/torch # 放在最顶部 import torch from torchvision import models weights models.VGG16_Weights.IMAGENET1K_V1 model models.vgg16(weightsweights) model.eval() print(模型加载完成)第一次运行会显示下载进度条下载完成后可以检查文件ls -lh /data/models/torch/checkpoints/你会看到vgg16-397923af.pth体积约528MB。第二次再运行脚本就没有下载进度条了说明缓存命中。这里多说一句为什么推荐weightsIMAGENET1K_V1而不是直接写weightsDEFAULT。虽然对于VGG16来说两者指向同一个权重但显式写出具体的IMAGENET1K_V1等于把我用的到底是哪一套权重写进了代码。以后torchvision如果上线了质量更好的新权重DEFAULT可能会跟着变而你的显式声明不会悄悄改变行为。复现实验的时候这种明确的记录非常重要。3.3 用weights自带的预处理做一次真实推理加载模型只是第一步。VGG16在ImageNet上训练时输入图像有一套固定的预处理流程先缩放到256再中心裁剪到224然后归一化。手工记这些数值容易出错好在weights.transforms()能直接帮你生成这套流程。from PIL import Image from torchvision import models weights models.VGG16_Weights.IMAGENET1K_V1 model models.vgg16(weightsweights) model.eval() preprocess weights.transforms() # 自动生成预处理 img Image.open(demo.jpg).convert(RGB) x preprocess(img).unsqueeze(0) # 加一个batch维度 with torch.no_grad(): logits model(x) pred logits.argmax(dim1).item() label weights.meta[categories][pred] print(预测类别:, label)weights.meta[categories]是权重自带的ImageNet类别列表一共1000类可以直接用来对应预测索引。如果只是想验证模型能不能跑通用torch.randn(1, 3, 224, 224)造一个随机张量喂进去输出形状是torch.Size([1, 1000])也足够确认加载成功了。4. 下载与加载的踩坑清单4.1 下到一半断了残留文件导致的诡异报错预训练模型文件动辄几百MB网络稍有波动就可能下载中断。新版PyTorch在下载时会先写临时文件下载完成后再改名成正式文件但你如果中途把进程杀掉或是磁盘空间不够仍然可能留下不完整的临时文件甚至损坏的.pth。典型症状是下一次运行报EOFError: Ran out of input或者torch.load时突然卡在奇怪的解序列化错误。这种时候先别怀疑代码去缓存目录看看有没有零字节或明显偏小的文件把它们删掉重新下载。如果你在大文件下载上经常翻车我习惯用wget -c做断点续传下完再手动放到目标目录全程自己掌控wget -c https://download.pytorch.org/models/vgg16-397923af.pth -O /data/models/torch/checkpoints/vgg16-397923af.pth下载完之后还可以用torch.hub.load_state_dict_from_url(..., check_hashTrue)做一次哈希校验确认文件内容和官方一致。比如直接运行import torch torch.hub.load_state_dict_from_url( https://download.pytorch.org/models/vgg16-397923af.pth, model_dir/data/models/torch/checkpoints, check_hashTrue )如果校验不通过它会明确告诉你下载的文件损坏而不是等到加载时才报一个看不懂的错误。4.2 手动改名导致重复下载前面说过缓存识别靠的是完整文件名。很多人手动下载后喜欢改成vgg16.pth这种好记的名字结果放回缓存目录后PyTorch发现vgg16-397923af.pth不存在又默默重新下载了一遍。磁盘上同时躺着两个五百多兆的文件问题还不大但下载时间浪费了。所以规则很简单如果你走缓存机制文件名必须保持原样如果你想自定义文件名就走方案二的torch.load手动加载别两边混着来。4.3 多用户共享目录的权限与抢占问题团队共用一台服务器大家约定把权重放在/data/models/torch。可一旦有两个用户同时第一次运行vgg16(weights...)就会同时触发下载往同一个文件路径写数据。轻则互相覆盖、下载重复重则出现权限冲突或临时文件残留。解决思路是让下载这件事只发生一次。我见过比较稳妥的做法是由管理员或第一个用户把权重下载好之后把目录权限改成只读或者准备一个预下载脚本明确告诉所有人新模型先跑这个脚本下载不要在代码里触发。这样平时大家跑代码都只读缓存既不会重复下载也不会出现权限打架。顺带一提共享目录建议放在独立挂载的存储上别放系统盘。VGG16一个文件五百多MB科研机器上ResNet、DenseNet再来几个系统盘很容易被撑爆影响整机稳定性。4.4 新旧API混用带来的坑torchvision从0.13开始推荐用weights参数替代早期的pretrained参数。如果你看到类似pretrainedis deprecated and will be removed的警告说明你的代码还在走老路子。老API写法model models.vgg16(pretrainedTrue)新API写法model models.vgg16(weightsmodels.VGG16_Weights.IMAGENET1K_V1)两者最终下载的是同一个权重文件但新API更好用它能显式声明权重版本、提供transforms()预处理、附带meta信息。更要注意的是pretrained和weights不要同时传否则会直接报参数冲突。另外weightsNone表示不加载预训练权重模型参数是随机初始化的做迁移学习微调时记得要明确区分从头训练和加载预训练这两种意图。5. 离线机器和团队协作怎么落地5.1 一次下载多机复用的搬运流程离线机器没有外网为什么还是推荐用models.vgg16(weights...)这种写法因为只要文件到位PyTorch根本不会去请求网络它只认缓存目录里有没有对应的文件。完整的搬运流程是这样的在一台有网的机器上设置好TORCH_HOME运行一次下载脚本确保$TORCH_HOME/checkpoints下出现目标权重文件。把整个checkpoints目录或者只需要的那几个.pth打包拷到离线机器。在离线机器上设置相同的TORCH_HOME把文件放到$TORCH_HOME/checkpoints下。正常运行models.vgg16(weights...)不再需要外网。拷贝之前建议顺手跑一下md5sum和官方文件名里那串哈希对照一下确认文件完整。文件名前缀的哈希本来就是MD5的前8位这个校验方法我一直觉得是PyTorch给的一个很贴心的设计。5.2 团队标准目录与自动化脚本如果要管理的不只一个模型强烈建议弄一个标准目录比如/data/pretrained/torch/checkpoints/然后写一个预下载脚本把需要的模型URL一次性下好。这里提供一个VGG系列的参考脚本#!/usr/bin/env bash set -e TORCH_HOME/data/pretrained/torch mkdir -p $TORCH_HOME/checkpoints cd $TORCH_HOME/checkpoints URLS( https://download.pytorch.org/models/vgg11-bbd30ac9.pth https://download.pytorch.org/models/vgg13-c768596a.pth https://download.pytorch.org/models/vgg16-397923af.pth https://download.pytorch.org/models/vgg19-dcbb9e9d.pth ) for url in ${URLS[]}; do echo downloading $(basename $url) wget -c -q $url done至于其他模型的URL可以直接用Python打印不用反复翻文档from torchvision.models import VGG16_Weights print(VGG16_Weights.IMAGENET1K_V1.url)拿ResNet举例把VGG16_Weights换成ResNet50_Weights同样能拿到对应URLtorchvision里所有带预训练权重的模型都支持这个操作。在Docker部署场景可以把权重文件直接打进镜像或者挂载一个共享卷。容器内通过启动脚本统一设置TORCH_HOME这样不管在哪台机器上启动模型加载行为都一致不会出现这台机器能跑换台机器就重新下载的尴尬。我自己现在养成的习惯是每接一个新的视觉项目第一件事就是确定权重存放根目录然后把需要的模型URL清单提交到项目仓库里。这样无论是本地调试、服务器训练还是离线部署都只有一份权威的模型来源不会各下各的。至于下载加载本身理解了缓存机制和文件名规则之后剩下的都是水到渠成的事。