我来帮你分析下这个问题：你的RTX 3090虽然有24GB显存，但UNET在处理1024×1360这种大尺寸图像时，中间卷积层的特征图会占用大量显存——哪怕batch size设为1，单张图对应的多层特征图加起来也很容易突破显存上限。缩小数据集后随机停止的问题，大概率是显存碎片化或者数据加载环节有异常，下面是按优先级排序的解决方案：

• 优先缩小输入图像尺寸
  1024×1360的分辨率对于UNET来说实在太大了，建议直接把图像和掩码缩放到原尺寸的1/2（512×680），或者训练时采用随机裁剪（比如每次从原图中裁剪出512×512的区域）。这样单张图的显存占用会直接降到原来的1/4，能快速解决OOM问题。如果担心精度损失，可以训练稳定后再尝试逐步增大尺寸。

• 优化TensorFlow显存分配策略
  默认情况下TensorFlow会抢占所有可用显存，容易导致其他程序或系统进程占用的显存和训练冲突。你可以设置动态显存增长，让TF按需申请显存：

  import tensorflow as tf
  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
          print(f"成功启用动态显存分配，检测到{len(gpus)}个物理GPU")
      except RuntimeError as e:
          print(f"显存设置失败: {e}")
  

  要是动态分配还是有问题，也可以直接限制显存上限（比如给训练分配20GB，留4GB给系统）：

  tf.config.set_logical_device_configuration(
      gpus[0],
      [tf.config.LogicalDeviceConfiguration(memory_limit=20480)]
  )
  

• 模型轻量化改造
  标准UNET的通道数设置（比如第一层64个通道）对大尺寸图像来说太耗显存，你可以做这些调整：

  • 减少卷积层的通道数：把编码器第一层的通道数从64改成32，后续层按比例减半（比如64→32，128→64等），先保证能训练起来，再逐步调优。
  • 替换为深度可分离卷积：用tf.keras.layers.DepthwiseConv2D替代普通的Conv2D，能大幅减少参数数量和显存占用，同时保持模型的特征提取能力。
  • 启用混合精度训练：让TF自动用float16执行部分运算，显存占用直接减半，且几乎不损失精度：

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

    注意最后输出层要确保是float32（如果掩码是float32的话），可以在输出层加dtype='float32'。

• 优化数据加载流程
  避免一次性把所有大尺寸图像加载到内存，改用tf.data.Dataset按需加载并预处理，同时在加载阶段就完成缩放/裁剪，减少内存压力：

  def load_and_preprocess(image_path, mask_path):
      # 读取图像并缩放
      img = tf.io.read_file(image_path)
      img = tf.image.decode_jpeg(img, channels=3)
      img = tf.image.resize(img, (512, 680))
      img = tf.cast(img, tf.float32) / 255.0
      
      # 读取掩码（用最近邻插值避免模糊）
      mask = tf.io.read_file(mask_path)
      mask = tf.image.decode_png(mask, channels=1)
      mask = tf.image.resize(mask, (512, 680), method='nearest')
      mask = tf.cast(mask, tf.float32) / 255.0
      
      return img, mask
  
  # 构建数据集
  image_paths = [你的图像路径列表]
  mask_paths = [对应的掩码路径列表]
  dataset = tf.data.Dataset.from_tensor_slices((image_paths, mask_paths))
  dataset = dataset.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
  dataset = dataset.batch(1).prefetch(tf.data.AUTOTUNE)
  

  这样训练时只会加载当前batch的图像，不会占用过多内存。

• 解决缩小数据集后随机停止的问题
  这个问题大概率是数据损坏或者显存碎片化：

  • 检查所有图像和掩码文件，看有没有损坏的（比如无法正常解码的文件），可以在加载函数中加入异常捕获。
  • 如果是显存碎片化，建议先按前面的方法缩小图像尺寸，再配合动态显存分配，基本能解决随机停止的问题。

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