Detectron2训练报错：RuntimeError输入与权重张量类型不匹配

问题 #

在搭载CUDA的Ubuntu笔记本本地Jupyter Notebook中运行基于Detectron2的图像检测代码时，出现如下报错：

RuntimeError: Input type (torch.cuda.DoubleTensor) and weight type (torch.cuda.FloatTensor) should be the same

输入数据为14张JPG图片及COCO格式标注文件，数据集注册环节无异常。

相关代码 #

dataset_dicts = load_coco_json("../Downsampled/Train/annotations.json", "../Downsampled/Train", dataset_name="Seals", extra_annotation_keys=None)
register_coco_instances("seal", {}, "../Downsampled/Train/annotations.json", "../Downsampled/Train")

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
#cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_50_FPN_3x.yaml"))
cfg.INPUT.MAX_SIZE_TRAIN = 2000
cfg.INPUT.MIN_SIZE_TRAIN = 800
cfg.INPUT.MAX_SIZE_TEST = 5000
cfg.INPUT.MIN_SIZE_TEST = 1000
cfg.DATASETS.TRAIN = ("seal",)
cfg.MODEL.PIXEL_MEAN = [mean_pixel_values[2], mean_pixel_values[1], mean_pixel_values[0]]
cfg.MODEL.PIXEL_STD = [std_pixel_values[2], std_pixel_values[1], std_pixel_values[0]]
cfg.DATASETS.TEST = ()  # no metrics implemented for this dataset
#cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.DEVICE="cuda"
cfg.MODEL.WEIGHTS = "detectron2://ImageNetPretrained/MSRA/R-50.pkl"  # initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.01
cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[8,16,32]]
cfg.MODEL.ANCHOR_GENERATOR.ANGLES = [[0,90]]
cfg.SOLVER.MAX_ITER = (
    300
)  # 300 iterations seems good enough, but you can certainly train longer
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (
    16
)  # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4  # 7 classes (none, error, pup, adult male, subadult male, female, juvenile)

print("Start training")
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

完整报错栈 #

ERROR [08/09 01:39:54 d2.engine.train_loop]: Exception during training:
Traceback (most recent call last):
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/engine/train_loop.py", line 149, in train
    self.run_step()
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/engine/defaults.py", line 494, in run_step
    self._trainer.run_step()
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/engine/train_loop.py", line 274, in run_step
    loss_dict = self.model(data)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 158, in forward
    features = self.backbone(images.tensor)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/modeling/backbone/fpn.py", line 139, in forward
    bottom_up_features = self.bottom_up(x)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/modeling/backbone/resnet.py", line 445, in forward
    x = self.stem(x)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/modeling/backbone/resnet.py", line 356, in forward
    x = self.conv1(x)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/hysterio/.pyenv/versions/3.7.6/envs/vivadata/lib/python3.7/site-packages/detectron2/layers/wrappers.py", line 114, in forward
    x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups
RuntimeError: Input type (torch.cuda.DoubleTensor) and weight type (torch.cuda.FloatTensor) should be the same
[08/09 01:39:54 d2.engine.hooks]: Total training time: 0:00:00 (0:00:00 on hooks)
[08/09 01:39:55 d2.utils.events]:  iter: 0    lr: N/A  max_mem: 249M

解决方法 #

报错核心是输入数据的张量类型（DoubleTensor）和模型权重的张量类型（FloatTensor）不匹配，Detectron2预训练模型默认用Float32精度，而输入数据被处理成了Float64（Double）精度。

具体修复步骤 #

1. 转换均值和方差类型
   你计算的mean_pixel_values和std_pixel_values大概率是float64类型，将其转为float32即可：

mean_pixel_values = mean_pixel_values.astype('float32')
std_pixel_values = std_pixel_values.astype('float32')

2. 自定义数据加载器强制转类型
   如果上述方法无效，自定义DatasetMapper确保输入张量为Float32：

from detectron2.data import DatasetMapper, build_detection_train_loader

class CustomDatasetMapper(DatasetMapper):
    def __call__(self, dataset_dict):
        dataset_dict = super().__call__(dataset_dict)
        dataset_dict['image'] = dataset_dict['image'].to(torch.float32)
        return dataset_dict

# 替换训练器的默认数据加载器
trainer = DefaultTrainer(cfg)
trainer.build_train_loader = lambda cfg: build_detection_train_loader(cfg, mapper=CustomDatasetMapper(cfg, True))

3. 全局设置默认张量类型
   在代码开头添加以下代码，强制PyTorch默认使用float32：

import torch
torch.set_default_dtype(torch.float32)

推荐方案 #

优先用第一种方法，因为问题基本源于均值和方差的类型不匹配，导致后续数据预处理环节将图像张量转为double类型，和模型float32权重冲突。

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