在Google Colab运行YOLO-NAS预测时遇类型不匹配RuntimeError求助

问题场景 #

使用Windows 10系统，在Google Colab中基于YOLO-NAS自定义数据集教程运行预测，调用最优模型时触发如下错误：
RuntimeError: Input type (unsigned char) and bias type (c10::Half) should be the same

完整错误栈：

RuntimeError                              Traceback (most recent call last)
<ipython-input-14-3169f5cc4e6d> in <cell line: 11>()
      9 
     10 test_image = 'test.JPG'
---> 11 best_model.predict(test_image).show()

16 frames
/usr/local/lib/python3.10/dist-packages/super_gradients/training/models/detection_models/customizable_detector.py in predict(self, images, iou, conf, fuse_model)
    175         """
    176         pipeline = self._get_pipeline(iou=iou, conf=conf, fuse_model=fuse_model)
---> 177         return pipeline(images)  # type: ignore
    178 
    179     def predict_webcam(self, iou: Optional[float] = None, conf: Optional[float] = None, fuse_model: bool = True):

/usr/local/lib/python3.10/dist-packages/super_gradients/training/pipelines/pipelines.py in __call__(self, inputs, batch_size)
     94             return self.predict_video(inputs, batch_size)
     95         elif check_image_typing(inputs):
---> 96             return self.predict_images(inputs, batch_size)
     97         else:
     98             raise ValueError(f"Input {inputs} not supported for prediction.")

/usr/local/lib/python3.10/dist-packages/super_gradients/training/pipelines/pipelines.py in predict_images(self, images, batch_size)
    109         images = load_images(images)
    110         result_generator = self._generate_prediction_result(images=images, batch_size=batch_size)
---> 111         return self._combine_image_prediction_to_images(result_generator, n_images=len(images))
    112 
    113     def predict_video(self, video_path: str, batch_size: Optional[int] = 32) -> VideoPredictions:

/usr/local/lib/python3.10/dist-packages/super_gradients/training/pipelines/pipelines.py in _combine_image_prediction_to_images(self, images_predictions, n_images)
    288         if n_images is not None and n_images == 1:
    289             # Do not show tqdm progress bar if there is only one image
---> 290             images_predictions = [next(iter(images_predictions))]
    291         else:
    292             images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Images")]

/usr/local/lib/python3.10/dist-packages/super_gradients/training/pipelines/pipelines.py in _generate_prediction_result(self, images, batch_size)
    147         else:
    148             for batch_images in generate_batch(images, batch_size):
---> 149                 yield from self._generate_prediction_result_single_batch(batch_images)
    150 
    151     def _generate_prediction_result_single_batch(self, images: Iterable[np.ndarray]) -> Iterable[ImagePrediction]:

/usr/local/lib/python3.10/dist-packages/super_gradients/training/pipelines/pipelines.py in _generate_prediction_result_single_batch(self, images)
    174             if self.fuse_model:
    175                 self._fuse_model(torch_inputs)
---> 176             model_output = self.model(torch_inputs)
    177             predictions = self._decode_model_output(model_output, model_input=torch_inputs)
    178 

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.10/dist-packages/super_gradients/training/models/detection_models/customizable_detector.py in forward(self, x)
     85 
     86     def forward(self, x):
---> 87         x = self.backbone(x)
     88         x = self.neck(x)
     89         return self.heads(x)

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.10/dist-packages/super_gradients/modules/detection_modules.py in forward(self, x)
     78         all_layers = ["stem"] + [f"stage{i}" for i in range(1, self.num_stages + 1)] + ["context_module"]
     79         for layer in all_layers:
---> 80             x = getattr(self, layer)(x)
     81             if layer in self.out_layers:
     82                 outputs.append(x)

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.10/dist-packages/super_gradients/training/models/detection_models/yolo_nas/yolo_stages.py in forward(self, x)
    136 
    137     def forward(self, x: Tensor) -> Tensor:
---> 138         return self.conv(x)
    139 
    140 

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.10/dist-packages/super_gradients/modules/qarepvgg_block.py in forward(self, inputs)
    177     def forward(self, inputs):
    178         if self.fully_fused:
---> 179             return self.se(self.nonlinearity(self.rbr_reparam(inputs)))
    180 
    181         if self.partially_fused:

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in forward(self, input)
    461 
    462     def forward(self, input: Tensor) -> Tensor:
---> 463         return self._conv_forward(input, self.weight, self.bias)
    464 
    465 class Conv3d(_ConvNd):

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    457                             weight, bias, self.stride,
    458                             _pair(0), self.dilation, self.groups)
---> 459         return F.conv2d(input, weight, bias, self.stride,
    460                         self.padding, self.dilation, self.groups)
    461 

RuntimeError: Input type (unsigned char) and bias type (c10::Half) should be the same

解决建议 #

• 强制模型转为float32类型
  加载最优模型后，立即将模型参数转换为float32，避免半精度与输入类型冲突：

  best_model = best_model.float()
  

• 手动处理输入图像的类型
  跳过默认的predict方法的自动加载，手动将图像转为float32张量并归一化：

  from PIL import Image
  import numpy as np
  import torch
  
  test_image = Image.open('test.JPG').convert('RGB')
  # 转换为模型所需的张量格式：(batch, channels, height, width)，并转为float32
  input_tensor = torch.tensor(np.array(test_image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
  # 如果使用GPU，将张量移到GPU
  if torch.cuda.is_available():
      input_tensor = input_tensor.cuda()
  # 直接前向传播获取预测结果
  predictions = best_model(input_tensor)
  # 可使用SuperGradients的可视化工具展示结果
  

• 关闭模型融合功能
  调用predict时指定fuse_model=False，避免融合过程中出现类型不兼容：

  best_model.predict(test_image, fuse_model=False).show()
  

• 检查模型加载配置
  若从 checkpoint 加载模型，确保加载时未强制半精度，必要时手动转换类型：

  best_model.load_state_dict(torch.load('best_model.pth'), strict=False)
  # 转为float32并移到GPU（如果可用）
  best_model = best_model.float()
  if torch.cuda.is_available():
      best_model = best_model.cuda()
  

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