TensorFlow Estimator CPU训练正常，GPU训练遇NaN及内存占用过高问题

一、GPU内存占满是否正常？ #

TensorFlow 1.x默认会预分配全部GPU内存，哪怕实际训练用不到这么多，所以你看到V100的32GB内存被占满是默认行为，不属于异常，但可以通过配置调整为按需分配。

二、降低GPU内存占用的方法 #

• 开启动态内存分配：这是最有效的方式，让TensorFlow只申请实际需要的内存，代码如下：

import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
tf.keras.backend.set_session(sess)  # 用Keras封装模型时需要绑定session

• 指定单GPU训练：如果机器有多块GPU，确保只用到目标V100，避免其他GPU抢占资源：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 根据nvidia-smi输出调整设备编号

• 清理无关进程：用nvidia-smi命令查看GPU进程，关闭其他训练任务、闲置的jupyter内核等占用GPU的进程。

三、解决PredictedDepth的NaN错误 #

NaN基本是数值不稳定导致的，针对深度估计模型，按以下步骤排查：

• 改用逆深度计算损失：直接用原始深度容易出现超大值导致梯度爆炸，把深度预测转为1/depth（逆深度），损失计算基于逆深度，能有效避免数值溢出。
• 添加梯度裁剪：限制梯度的最大范数，防止梯度爆炸导致参数更新异常：

optimizer = tf.train.AdamOptimizer(1e-4)
grads_vars = optimizer.compute_gradients(total_loss)
# 把梯度范数限制在1.0以内，可根据实际情况调整
clipped_grads_vars = [(tf.clip_by_norm(g, 1.0), v) for g, v in grads_vars if g is not None]
train_op = optimizer.apply_gradients(clipped_grads_vars)

• 检查输入数据：确保训练图像的像素值已正确归一化（比如到[0,1]或[-1,1]），排查是否存在损坏的样本（像素值为NaN或无穷大）。
• 权重初始化优化：避免用过大的初始权重，改用方差缩放初始化：

initializer = tf.variance_scaling_initializer(scale=2.0, mode='fan_in', distribution='truncated_normal')

• 开启数值调试：在预测深度输出后添加检查，定位NaN出现的具体位置：

predicted_depth = ...  # 你的深度预测张量
tf.debugging.check_numerics(predicted_depth, "Predicted depth contains NaN/Inf")

四、额外注意事项 #

• 确认CUDA和cuDNN版本匹配：TensorFlow 1.15要求CUDA10.0搭配cuDNN7.4.x，版本不兼容可能引发内存或数值异常。
• 清理计算图：训练前用tf.reset_default_graph()清理残留的图操作，避免无用张量占用内存。

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