Google Colab升级TensorFlow 2.7后OOM报错及张量形状异常问题

该问题是TensorFlow 2.7版本Keras组件的兼容性bug导致：Keras官方Zero-DCE示例中所有Concatenate层未显式指定拼接轴，依赖默认参数实现通道维度拼接，TensorFlow 2.7调整了Concatenate层的默认轴判断逻辑，误将通道维拼接改为batch后的第二维拼接，原本应输出的(16,1024,1024,64)张量被错误生成为(16,64,1024,1024)，显存占用直接翻64倍触发OOM。
你尝试降级TensorFlow出现的CuDNN不兼容问题，是因为Colab当前预装的CUDA版本与TensorFlow 2.6要求的CUDA 11.2不匹配导致。

• 找到代码中所有Concatenate层的调用位置，手动添加axis=-1参数，强制指定按最后一维（通道维）拼接，示例修改如下：
  原代码：x = Concatenate()([x1, x2])
  修改后：x = Concatenate(axis=-1)([x1, x2])
• 若修改后仍有轻微显存不足的情况，将批量大小batch_size从默认的16下调至8或4即可。
• 额外可在代码首段加入显存动态分配逻辑，避免无效显存占用：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

执行以下命令安装匹配依赖，执行完后重启Colab运行时即可正常使用TensorFlow 2.6：

!apt-get update -q
!apt-get install cuda-11-2 -q
!pip install tensorflow==2.6.5 tensorflow-gpu==2.6.5

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