使用mmsegmentation训练自定义数据集时遇张量尺寸不匹配RuntimeError

我正在使用MMSegmentation对自定义数据集执行语义分割任务，已编写专属配置文件。训练过程中触发RuntimeError，提示批量内张量尺寸不统一。我已确认数据集中所有图片尺寸一致，但无法定位程序中触发报错的具体代码，因此不知如何解决。

报错堆栈信息 #

File "tools/train.py", line 104, in <module>
    main()
  File "tools/train.py", line 100, in main
    runner.train()
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1721, in train
    model = self.train_loop.run()  # type: ignore
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 278, in run
    self.run_iter(data_batch)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 301, in run_iter
    outputs = self.runner.model.train_step(
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/wrappers/distributed.py", line 121, in train_step
    losses = self._run_forward(data, mode='loss')
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/wrappers/distributed.py", line 161, in _run_forward
    results = self(**data, mode=mode)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1040, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1000, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/segmentors/base.py", line 94, in forward
    return self.loss(inputs, data_samples)
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/segmentors/encoder_decoder.py", line 176, in loss
    loss_decode = self._decode_head_forward_train(x, data_samples)
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/segmentors/encoder_decoder.py", line 137, in _decode_head_forward_train
    loss_decode = self.decode_head.loss(inputs, data_samples,
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/decode_heads/decode_head.py", line 262, in loss
    losses = self.loss_by_feat(seg_logits, batch_data_samples)
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/decode_heads/decode_head.py", line 305, in loss_by_feat
    seg_label = self._stack_batch_gt(batch_data_samples)
  File "/home/deshpand/Thesis/semantic_segmentation_network/mmseg_new/mmsegmentation/mmseg/models/decode_heads/decode_head.py", line 289, in _stack_batch_gt
    return torch.stack(gt_semantic_segs, dim=0)
RuntimeError: stack操作要求每个张量尺寸一致，但第0个张量尺寸为[421, 604, 516]，第1个为[512, 513, 516]

...（重复报错堆栈省略，核心错误均为同批次内GT张量尺寸不一致）

ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 1820710) of binary: /home/deshpand/anaconda3/envs/openmmlab/bin/python
Traceback (most recent call last):
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/launch.py", line 195, in <module>
    main()
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/launch.py", line 191, in main
    launch(args)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/launch.py", line 176, in launch
    run(args)
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/run.py", line 753, in run
    elastic_launch(
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 132, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/deshpand/anaconda3/envs/openmmlab/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 246, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

解决方法 #

• 检查数据预处理管道

  • 确认配置文件中的img_scale是否为固定值，若使用了RandomResize这类随机缩放增强，必须配合Pad操作（设置固定尺寸或size_divisor），确保同批次所有图片和标签尺寸统一。
  • 若使用RandomCrop，需保证裁剪尺寸固定，且后续用Pad补全到统一尺寸，避免同批次样本尺寸差异。

• 验证标签数据一致性

  • 即使原始图片尺寸一致，标签可能存在导出或预处理错误，写脚本遍历所有标签文件，检查其尺寸是否与对应图片完全匹配。

• 调试定位问题样本

  • 在mmseg/models/decode_heads/decode_head.py的_stack_batch_gt方法中，加入打印语句输出每个gt_semantic_segs的尺寸，找到异常样本后回溯其预处理流程。
  • 暂时关闭分布式训练，改用单卡训练，报错信息会更简洁，便于定位问题。

• 配置文件修正示例
  如果使用随机增强，修改训练管道确保尺寸统一：

  train_pipeline = [
      dict(type='LoadImageFromFile'),
      dict(type='LoadAnnotations'),
      dict(type='RandomResize', scale=(1024, 1024), ratio_range=(0.5, 2.0)),
      dict(type='RandomCrop', crop_size=(512, 512), cat_max_ratio=0.75),
      # 新增Pad操作固定尺寸
      dict(type='Pad', size=(512, 512), pad_val=0, seg_pad_val=255),
      dict(type='RandomFlip', prob=0.5),
      dict(type='PackSegInputs')
  ]
  

  或者直接固定缩放尺寸，去掉随机缩放：

  train_pipeline = [
      dict(type='LoadImageFromFile'),
      dict(type='LoadAnnotations'),
      dict(type='Resize', scale=(512, 512), keep_ratio=True),
      dict(type='Pad', size_divisor=32),
      dict(type='RandomFlip', prob=0.5),
      dict(type='PackSegInputs')
  ]
  

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