【权重】离线环境怎么用预训练权重

Rabbit_QL2026-04-23 13:27

在真实的生产环境中，模型加载往往不会像教程里那么"顺滑"。

服务器无法访问外网（安全隔离环境）
内网集群（HPC / 企业训练平台）
Docker 镜像需要预置模型
推理服务必须离线启动

你会发现，torchvision 的"自动下载预训练权重"机制完全失效。这时候，就需要我们手动接管整个流程。

一、Step1: 在有网的机器上下好权重

python 复制代码

from torchvision import models
_ = models.vgg16(weights='IMAGENET1K_V1')   # 触发一次自动下载

检查本地 cache 是否存在权重，如果没有 → 自动从官方 URL 下载，然后缓存到 ~/.cache/torch/hub/checkpoints/。

但在离线环境，因为无法访问下载地址，自动流程直接失败，模型初始化被阻断。

二、Step2：定位权重文件

下载完成后，权重会被存到 torch cache 中。

python 复制代码

import torch, os
from torchvision import models

weights = models.VGG16_Weights.IMAGENET1K_V1
filename = weights.url.split('/')[-1]        # vgg16-397923af.pth
src = os.path.join(torch.hub.get_dir(), 'checkpoints', filename)

三、Step3复制到目标机器

scp / 放进镜像 / 拷到内网，然后在无网机器上加载：

python 复制代码

state_dict = torch.load('/path/to/vgg16-397923af.pth', map_location='cpu')
model = models.vgg16(weights=None)          # 注意是 None，别触发下载
model.load_state_dict(state_dict)

关键点 ：weights=None 告诉 torchvision "只给我结构、别下载"。然后你手动 load_state_dict。

这一套流程在 HuggingFace 的 transformers 里完全同构，学会一次到处能用。