首先直接给结论：你遇到的ResourceExhaustedError确实是GPU显存不足导致的，咱们算笔账就能一目了然——报错里提到的张量shape[50,32,512,512]是float32类型，每个元素占4字节，总显存需求是：
50 × 32 × 512 × 512 × 4 = 16,777,216,000 字节 ≈ 15.6GB
而GeForce 940MX的显存通常只有2GB或4GB，完全装不下这么大的张量，自然会触发显存分配失败。

至于CPU能正常运行，是因为CPU可以利用系统内存（甚至内存不够时会借硬盘虚拟内存），但这种方式速度极慢，所以单轮epoch要耗时6小时。

下面是针对你的情况的具体优化方案，按优先级排序：

1. 大幅降低Batch Size #

这是最直接有效的办法。你当前设置的batch_size=50对于940MX来说太大了，建议先改成2、4或者8测试能稳定运行的最大值。比如batch_size=2时，该张量的显存需求仅约671MB，加上UNet模型本身的参数（几十MB左右），2GB显存也能轻松应对。

修改代码里的这一行：

model.fit(imgs_train, imgs_mask_train, batch_size=2, epochs=10, verbose=1, shuffle=True, callbacks=[model_checkpoint])

2. 缩小输入图像尺寸 #

当前你用的是512×512的输入，如果改成256×256，每个样本的显存占用会降到原来的1/4，同样batch size下的显存压力会大幅降低。不过要注意，缩小图像可能会丢失一些结节的细节，对检测精度有一定影响，需要根据你的需求权衡。

修改代码里的两个尺寸参数：

img_rows = 256
img_cols = 256

3. 优化模型与显存分配策略 #

• 减少卷积通道数：把UNet里的通道数减半，比如conv1的32改成16，conv2的64改成32，以此类推。这样每个卷积层的特征图显存占用都会减少，同时模型参数也会更少，训练速度更快。
• 启用GPU内存增长模式：让TensorFlow按需分配显存，而不是一开始就占用全部显存，避免浪费。在代码开头添加：

  import tensorflow as tf
  tf.keras.backend.clear_session()
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

4. 启用混合精度训练 #

用float16类型存储部分张量，能把显存占用减少一半，同时几乎不影响模型精度。在Keras里可以通过以下方式配置：

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

940MX支持float16运算，完全可以使用这个方案。

额外建议 #

• 训练前用nvidia-smi命令（Windows下需要安装NVIDIA驱动的配套工具）检查GPU显存使用情况，确保没有其他进程占用显存。
• 如果想要近似大batch的训练效果但又没足够显存，可以试试梯度累积：用小batch训练，累积多个batch的梯度后再更新参数。比如用batch_size=2，累积25次再更新，效果近似于batch_size=50，但显存只需要小batch的量，不过需要手动修改代码处理梯度累积逻辑。

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