不依赖torchvision加载含架构和权重的PyTorch .pth模型是否可行？

你遇到的ModuleNotFoundError和torch的serialization.py本身的依赖无关：你手中的vgg.pth是直接序列化了完整的VGG模型实例存储的，torch.load底层依赖pickle做反序列化，重建实例时需要导入该模型类的原始定义，而这个VGG类属于torchvision模块，因此pickle加载时会自动尝试导入torchvision，触发缺失报错。

下面是三种可在不安装torchvision的前提下加载模型的方案，可根据你的场景选择：

方案1：自定义反序列化类，直接提取权重（推荐，适合仅需要权重的场景） #

自定义pickle的Unpickler拦截torchvision类的导入请求，不需要实现完整模型结构就能提取出state_dict：

import os
import torch
import pickle
from collections import OrderedDict

class CustomUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        # 拦截所有torchvision下的类加载请求，返回空占位类
        if module.startswith('torchvision'):
            return type(name, (), {})
        return super().find_class(module, name)

with open(r'.\vgg.pth', 'rb') as f:
    loaded_obj = CustomUnpickler(f).load()

# 提取权重参数
state_dict = loaded_obj.state_dict() if hasattr(loaded_obj, 'state_dict') else loaded_obj

提取到的state_dict可以直接加载到你自己实现的同结构VGG模型中。

方案2：模拟torchvision模块，加载完整模型实例 #

如果需要得到完整的可运行模型实例，你可以先自己实现和torchvision官方结构完全一致的VGG类，再模拟torchvision的模块路径供pickle查找：

import sys
import torch

# 自行实现的VGG类，结构必须和torchvision官方VGG完全匹配
class VGG(torch.nn.Module):
    def __init__(self, features, num_classes=1000, init_weights=True):
        super().__init__()
        self.features = features
        self.avgpool = torch.nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(512 * 7 * 7, 4096),
            torch.nn.ReLU(True),
            torch.nn.Dropout(),
            torch.nn.Linear(4096, 4096),
            torch.nn.ReLU(True),
            torch.nn.Dropout(),
            torch.nn.Linear(4096, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    torch.nn.init.constant_(m.bias, 0)
            elif isinstance(m, torch.nn.BatchNorm2d):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0)
            elif isinstance(m, torch.nn.Linear):
                torch.nn.init.normal_(m.weight, 0, 0.01)
                torch.nn.init.constant_(m.bias, 0)

# 模拟torchvision模块路径
sys.modules['torchvision'] = type('mock_torchvision', (), {})()
sys.modules['torchvision.models'] = type('mock_models', (), {})()
sys.modules['torchvision.models.vgg'] = type('mock_vgg', (), {'VGG': VGG, 'vgg16': VGG})()

# 直接加载模型
model = torch.load(r'.\vgg.pth')

方案3：跨环境转存纯权重文件（最稳妥） #

如果有可安装torchvision的临时环境，可先在该环境下将完整模型转存为仅含权重的文件，后续加载就不需要依赖torchvision：

1. 有torchvision的环境执行：

import torch
model = torch.load(r'.\vgg.pth')
torch.save(model.state_dict(), 'vgg_weights_only.pth')

2. 无torchvision的环境执行（需自行实现同结构VGG）：

import torch
from your_vgg_code import VGG
model = VGG(...) # 构造和原模型结构一致的实例
model.load_state_dict(torch.load('vgg_weights_only.pth'))

内容的提问来源于stack exchange，提问作者Jonas De Schouwer