Python3.6环境下无需Ultralytics直接加载YOLO11n分类.pt模型方案咨询

方案1：修复TorchScript转换精度问题 #

• 固定输入尺寸转换：用torch.jit.trace而非torch.jit.script，转换时指定和训练一致的固定输入尺寸（比如(1,3,224,224)），避免动态维度导致的精度丢失。示例代码：

  import torch
  # 加载模型并设置为评估模式
  model = torch.load('your_trained_model.pt')['model'].float().eval()
  # 生成匹配输入尺寸的dummy张量
  dummy_input = torch.randn(1, 3, 224, 224)
  # 追踪模型并保存
  traced_model = torch.jit.trace(model, dummy_input)
  traced_model.save('yolo11n_classifier_ts.pt')
  

• 严格匹配预处理逻辑：推理时的图像resize、归一化（均值/方差）、通道顺序必须和训练阶段完全一致，比如训练用RGB通道、/255归一化，推理就不能改成BGR或其他归一化方式。

方案2：手动重构YOLO11分类架构并加载权重 #

基于PyTorch1.13.x（支持Python3.6的最后一个PyTorch大版本）手动复刻YOLO11n分类网络，直接加载权重文件：

1. 先在支持Ultralytics的环境导出纯权重字典：

   from ultralytics import YOLO
   model = YOLO('your_trained_model.pt')
   torch.save(model.model.state_dict(), 'yolo11n_classifier_weights.pth')
   

2. 在Python3.6环境重构YOLO11n核心结构（参考Ultralytics开源的YOLO11代码，剥离Ultralytics依赖）：

   import torch
   import torch.nn as nn
   
   # 实现YOLO11基础模块
   class Conv(nn.Module):
       def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
           super().__init__()
           self.conv = nn.Conv2d(c1, c2, k, s, nn.modules.utils._pair(k//2) if p is None else p, groups=g, bias=False)
           self.bn = nn.BatchNorm2d(c2)
           self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
   
       def forward(self, x):
           return self.act(self.bn(self.conv(x)))
   
   class C2f(nn.Module):
       def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
           super().__init__()
           self.c = int(c2 * e)
           self.cv1 = Conv(c1, 2 * self.c, 1, 1)
           self.cv2 = Conv((2 + n) * self.c, c2, 1)
           self.m = nn.ModuleList(Conv(self.c, self.c, 3, 1, g=g) for _ in range(n))
   
       def forward(self, x):
           y = list(self.cv1(x).split((self.c, self.c), 1))
           y.extend(m(y[-1]) for m in self.m)
           return self.cv2(torch.cat(y, 1))
   
   class SPPF(nn.Module):
       def __init__(self, c1, c2, k=5):
           super().__init__()
           c_ = c1 // 2
           self.cv1 = Conv(c1, c_, 1, 1)
           self.cv2 = Conv(c_ * 4, c2, 1, 1)
           self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k//2)
   
       def forward(self, x):
           x = self.cv1(x)
           y1 = self.m(x)
           y2 = self.m(y1)
           return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
   
   # 构建完整YOLO11n分类模型
   class YOLO11nClassifier(nn.Module):
       def __init__(self, num_classes=10):
           super().__init__()
           # 骨干网络，严格匹配YOLO11n的层数和通道数
           self.backbone = nn.Sequential(
               Conv(3, 16, 3, 2),
               Conv(16, 32, 3, 2),
               C2f(32, 32, 1, True),
               Conv(32, 64, 3, 2),
               C2f(64, 64, 2, True),
               Conv(64, 128, 3, 2),
               C2f(128, 128, 2, True),
               Conv(128, 256, 3, 2),
               C2f(256, 256, 1, True),
               Conv(256, 512, 3, 2),
               C2f(512, 512, 1, True),
           )
           # 分类头
           self.head = nn.Sequential(
               SPPF(512, 512),
               Conv(512, 1024, 1),
               nn.AdaptiveAvgPool2d(1),
               nn.Flatten(),
               nn.Linear(1024, num_classes)
           )
   
       def forward(self, x):
           x = self.backbone(x)
           x = self.head(x)
           return x
   

3. 加载权重并推理：

   # 初始化模型，指定你的分类类别数
   model = YOLO11nClassifier(num_classes=你的类别数)
   # 加载权重
   model.load_state_dict(torch.load('yolo11n_classifier_weights.pth'))
   model.eval()
   
   # 预处理图像为张量（示例）
   def preprocess(img_path):
       from PIL import Image
       img = Image.open(img_path).resize((224,224)).convert('RGB')
       img = np.array(img)/255.0
       img = torch.tensor(img).permute(2,0,1).unsqueeze(0).float()
       return img
   
   # 推理
   with torch.no_grad():
       input_tensor = preprocess('test.jpg')
       outputs = model(input_tensor)
       pred_class = torch.argmax(outputs, dim=1).item()
   

方案3：转换为ONNX并使用ONNX Runtime推理 #

ONNX Runtime支持Python3.6，无需依赖PyTorch或Ultralytics：

1. 在支持Ultralytics的环境导出ONNX模型：

   from ultralytics import YOLO
   model = YOLO('your_trained_model.pt')
   # 指定固定输入尺寸，禁用动态维度
   model.export(format='onnx', imgsz=224, dynamic=False)
   

2. 在Python3.6环境安装兼容版本的ONNX Runtime：

   pip install onnxruntime==1.14.1
   

3. 编写推理脚本：

   import onnxruntime as ort
   import cv2
   import numpy as np
   
   # 加载ONNX模型
   sess = ort.InferenceSession('yolo11n_classifier.onnx')
   input_name = sess.get_inputs()[0].name
   output_name = sess.get_outputs()[0].name
   
   # 图像预处理（匹配训练逻辑）
   img = cv2.imread('test.jpg')
   img = cv2.resize(img, (224,224))
   img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
   img = img.transpose((2,0,1)).astype(np.float32)/255.0
   img = np.expand_dims(img, axis=0)
   
   # 推理
   outputs = sess.run([output_name], {input_name: img})
   pred_class = np.argmax(outputs[0], axis=1)[0]
   

内容的提问来源于stack exchange，提问作者Gayatri Chawda