Model.fit时出现broadcastable shapes错误的UNet适配2D MRI问题

问题背景 #

作为Python与机器学习新手，尝试将UNet代码适配到尺寸为(51251224)的2D MRI图像训练模型，使用的模型代码如下：

inputs = Input((512, 512, 24))
conv1 = Conv2D(32, 3, activation='relu', padding='same')(inputs)
conv1 = Conv2D(32, 3, activation='relu', padding='same')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

conv2 = Conv2D(64, 3, activation='relu', padding='same')(pool1)
conv2 = Conv2D(64, 3, activation='relu', padding='same')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

conv3 = Conv2D(128, 3, activation='relu', padding='same')(pool2)
conv3 = Conv2D(128, 3, activation='relu', padding='same')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

conv4 = Conv2D(256, 3, activation='relu', padding='same')(pool3)
conv4 = Conv2D(256, 3, activation='relu', padding='same')(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

conv5 = Conv2D(512, 3, activation='relu', padding='same')(pool4)
conv5 = Conv2D(512, 3, activation='relu', padding='same')(conv5)
    
up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=-1)
conv6 = Conv2D(256, 3, activation='relu', padding='same')(up6)
conv6 = Conv2D(256, 3, activation='relu', padding='same')(conv6)

up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=-1)
conv7 = Conv2D(128, 3, activation='relu', padding='same')(up7)
conv7 = Conv2D(128, 3, activation='relu', padding='same')(conv7)

up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=-1)
conv8 = Conv2D(64, 3, activation='relu', padding='same')(up8)
conv8 = Conv2D(64, 3, activation='relu', padding='same')(conv8)

up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=-1)
conv9 = Conv2D(32, 3, activation='relu', padding='same')(up9)
conv9 = Conv2D(32, 3, activation='relu', padding='same')(conv9)

conv10 = Conv2D(1, 1, activation='sigmoid')(conv9)

model = Model(inputs=[inputs], outputs=[conv10])

model.compile(optimizer=Adam(lr=1e-5),loss=dice_coef_loss, metrics=[dice_coef, 'accuracy'])

执行训练代码时：

model_checkpoint = ModelCheckpoint('weights2021.h5', monitor='val_loss', save_best_only=True)
history = model.fit(imgs_train, masks_train, batch_size=128, epochs=5, verbose=1, shuffle=True, validation_split=0.2, callbacks=[model_checkpoint])

出现如下错误：

Epoch 1/5
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-42-6b4fb6c38414> in <module>
      1 model_checkpoint = ModelCheckpoint('weights2021.h5', monitor='val_loss', save_best_only=True)
----> 2 history = model.fit(imgs_train, masks_train, batch_size=128, epochs=5, verbose=1, shuffle=True, validation_split=0.2, callbacks=[model_checkpoint])

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     53     ctx.ensure_initialized()
     54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---&gt; 55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:
     57     if name is not None:
InvalidArgumentError: Graph execution error:
Detected at node 'dice_coef_loss/mul' defined at (most recent call last):
    File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
      "__main__", mod_spec)
    File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py", line 16, in <module>
      app.launch_new_instance()
    File "/usr/local/lib/python3.7/dist-packages/traitlets/config/application.py", line 846, in launch_instance
      app.start()
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelapp.py", line 612, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.7/dist-packages/tornado/platform/asyncio.py", line 132, in start
      self.asyncio_loop.run_forever()
    File "/usr/lib/python3.7/asyncio/base_events.py", line 541, in run_forever
      self._run_once()
    File "/usr/lib/python3.7/asyncio/base_events.py", line 1786, in _run_once
      handle._run()
    File "/usr/lib/python3.7/asyncio/events.py", line 88, in _run
      self._context.run(self._callback, *self._args)
    File "/usr/local/lib/python3.7/dist-packages/tornado/ioloop.py", line 758, in _run_callback
      ret = callback()
    File "/usr/local/lib/python3.7/dist-packages/tornado/stack_context.py", line 300, in null_wrapper
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/tornado/gen.py", line 1233, in inner
      self.run()
    File "/usr/local/lib/python3.7/dist-packages/tornado/gen.py", line 1147, in run
      yielded = self.gen.send(value)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 365, in process_one
      yield gen.maybe_future(dispatch(*args))
    File "/usr/local/lib/python3.7/dist-packages/tornado/gen.py", line 326, in wrapper
      yielded = next(result)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 268, in dispatch_shell
      yield gen.maybe_future(handler(stream, idents, msg))
    File "/usr/local/lib/python3.7/dist-packages/tornado/gen.py", line 326, in wrapper
      yielded = next(result)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/kernelbase.py", line 545, in execute_request
      user_expressions, allow_stdin,
    File "/usr/local/lib/python3.7/dist-packages/tornado/gen.py", line 326, in wrapper
      yielded = next(result)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/ipkernel.py", line 306, in do_execute
      res = shell.run_cell(code, store_history=store_history, silent=silent)
    File "/usr/local/lib/python3.7/dist-packages/ipykernel/zmqshell.py", line 536, in run_cell
      return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2855, in run_cell
      raw_cell, store_history, silent, shell_futures)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 2881, in _run_cell
      return runner(coro)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/async_helpers.py", line 68, in _pseudo_sync_runner
      coro.send(None)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 3058, in run_cell_async
      interactivity=interactivity, compiler=compiler, result=result)
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 3249, in run_ast_nodes
      if (await self.run_code(code, result,  async_=asy)):
    File "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py", line 3326, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "<ipython-input-42-6b4fb6c38414>", line 2, in <module>
      history = model.fit(imgs_train, masks_train, batch_size=128, epochs=5, verbose=1, shuffle=True, validation_split=0.2, callbacks=[model_checkpoint])
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1384, in fit
      tmp_logs = self.train_function(iterator)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1021, in train_function
      return step_function(self, iterator)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1010, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1000, in run_step
      outputs = model.train_step(data)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 860, in train_step
      loss = self.compute_loss(x, y, y_pred, sample_weight)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 919, in compute_loss
      y, y_pred, sample_weight, regularization_losses=self.losses)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/compile_utils.py", line 201, in __call__
      loss_value = loss_obj(y_t, y_p, sample_weight=sw)
    File "/usr/local/lib/python3.7/dist-packages/keras/losses.py", line 141, in __call__
      losses = call_fn(y_true, y_pred)
    File "/usr/local/lib/python3.7/dist-packages/keras/losses.py", line 245, in call
      return ag_fn(y_true, y_pred, **self._fn_kwargs)
    File "<ipython-input-26-bfe16c112741>", line 11, in dice_coef_loss
      return -dice_coef(y_true, y_pred)
    File "<ipython-input-26-bfe16c112741>", line 6, in dice_coef
      intersection = K.sum(y_true_f * y_pred_f)
Node: 'dice_coef_loss/mul'
required broadcastable shapes
     [[{{node dice_coef_loss/mul}}]] [Op:__inference_train_function_28414]

使用的Dice损失函数代码：

smooth = 1
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return K.mean(2. * intersection) / (K.sum(y_true_f + y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

错误原因分析 #

1. 形状不匹配：错误日志提示required broadcastable shapes，说明y_true（真实mask）和y_pred（模型输出）的形状无法进行逐元素相乘操作。模型输出形状为(batch_size, 512, 512, 1)，但你的masks_train可能是(batch_size, 512, 512)（缺少通道维度）或(batch_size, 512, 512, 24)（通道数与输出不匹配）。
2. Dice系数计算错误：原函数中K.mean(2. * intersection)是多余操作，intersection已经是全局求和后的标量，K.mean不会改变其值；同时分子未添加smooth，可能导致极端情况下除零错误；使用负Dice作为损失虽然可行，但更常规的做法是用1 - dice_coef，优化方向更直观。
3. Batch Size过大：batch_size=128对于512x512x24的输入来说，会占用极大显存，容易引发内存溢出或计算异常。

解决步骤 #

1. 统一Mask形状 #

确保masks_train的形状与模型输出一致：

• 如果是单类别分割，mask应为(batch_size, 512, 512, 1)。如果当前mask是(batch_size, 512, 512)，添加通道维度：

  import numpy as np
  masks_train = np.expand_dims(masks_train, axis=-1)
  

• 如果是多类别分割（24类），需要修改模型最后一层的输出通道数：

  conv10 = Conv2D(24, 1, activation='sigmoid')(conv9)  # 或使用softmax激活，根据任务调整
  

2. 修正Dice损失函数 #

调整Dice系数计算逻辑，修复多余操作并添加平滑项：

smooth = 1e-5  # 使用更小的平滑值避免影响结果
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    # 分子分母都加smooth，避免除零，同时去掉多余的K.mean
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return 1 - dice_coef(y_true, y_pred)  # 常规损失形式，优化目标更清晰

3. 调整Batch Size #

将batch_size修改为显存可承受的范围，比如8或16：

history = model.fit(imgs_train, masks_train, batch_size=8, epochs=5, verbose=1, shuffle=True, validation_split=0.2, callbacks=[model_checkpoint])

4. 可选：验证输入输出形状 #

在训练前添加形状验证代码，确保数据与模型匹配：

print(f"Input shape: {imgs_train.shape}")
print(f"Mask shape: {masks_train.shape}")
print(f"Model output shape: {model.output_shape}")

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