Keras卷积红区检查内存报错，如何用Tesla V100高效运行大批次训练

• 调整输入数据精度
  你当前数据生成器中定义的x、y均为np.int64类型，对于绝大多数序列分类任务，输入特征用np.float32、标签用np.int32即可满足精度要求，仅这一项就能将输入数据的显存占用减半。如果你的输入特征取值范围较小（如0~255），甚至可以先用np.uint8加载，送入模型前再自动转换为浮点型，进一步降低CPU-GPU传输开销和显存占用。

• 移除模型冗余计算
  你当前的卷积层定义存在重复激活操作：Conv1D层已设置activation='relu'，后续又追加了BatchNormalization和一次Activation('relu')，既浪费算力也会额外存储中间激活值。标准的卷积-BN-激活流程应去掉Conv1D中的activation参数，仅保留BN后的激活层即可。

• 开启显存动态分配
  深度学习框架默认会预占几乎全部GPU显存，容易产生显存碎片导致分配失败，可开启按需分配模式：
  TensorFlow版本：

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

PyTorch版本：

import torch
torch.cuda.set_per_process_memory_fraction(0.95, 0)

• 启用混合精度训练
  Tesla V100支持Tensor Core加速，开启混合精度（float16与float32混合计算）可将显存占用降低近一半，同时训练速度提升30%~100%，无需修改模型结构：
  TensorFlow版本：

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

PyTorch版本可直接调用torch.cuda.amp的GradScaler工具实现。

• 关闭cudnn调试检查
  你遇到的报错是cudnn的越界检查功能占用额外显存导致，不影响模型正确性，可直接设置环境变量关闭该功能：

import os
os.environ['TF_CUDNN_DISABLE_CONV_REDZONE'] = '1'
# PyTorch环境设置为os.environ['CUDNN_CONV_REDZONE_CHECK'] = '0'

• 梯度累积等效扩大batch size
  如果完成上述优化后仍无法达到目标batch size，可使用梯度累积策略：比如目标等效batch size为20000，当前单卡最大可跑batch size为2000，就每跑10个batch再更新一次梯度，训练效果和直接用batch size 20000完全一致，无额外显存开销。

完成以上优化后，16GB V100可稳定运行batch size 8000~10000，搭配梯度累积可实现等效20000的batch size，400万样本的训练速度相比原batch size 1000的配置提升8~10倍。

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