为什么TensorFlow创建简单模型时会占用大量VRAM显存？

你的猜想基本是正确的，运行那两行极简模型代码时出现的4000+MiB显存占用，和你定义的单层全连接模型几乎没有关系，绝大部分是GPU计算栈初始化带来的固定开销，具体VRAM中加载的内容分为以下几部分：

• CUDA系列基础计算库：这部分占比最高，包括CUDA Runtime运行时、cuBLAS线性代数库、cuDNN深度学习加速库、cuFFT信号处理库等TensorFlow GPU运行的必备依赖，这些库加载到显存本身就要占用1.5~2.5GiB左右的空间，属于只要用GPU跑TensorFlow就必须付出的固定开销，和模型大小完全无关。
• TensorFlow GPU执行上下文：第一次触发GPU设备检测时，TensorFlow会初始化整套GPU调度环境，包括计算流管理单元、内核函数注册缓存、内存管理元数据、设备通信上下文，这部分通常要占数百MB到1GiB不等的空间。
• 默认策略的显存预占空间：TensorFlow旧版本默认启动时直接占满所有空闲GPU显存，新版本也会默认预留大块显存入自己的内存管理池，减少后续训练/推理时动态申请显存的开销。你在nvidia-smi中看到的占用数值，包含了这部分被TensorFlow提前划走、暂时还没实际使用的预留空间，不是模型真实消耗的显存。
• 模型本身的参数占用：你定义的Dense(1, input_shape=(1,))层总共只有2个FP32参数（1个权重、1个偏置），实际显存占用不到100字节，完全可以忽略。

如果想验证这个结论，可以在导入TensorFlow后开启显存按需增长模式，再定义模型查看显存占用：

import tensorflow as tf
# 开启GPU显存按需分配，不提前预占大块空间
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

开启该配置后，初始显存占用会降到1GiB左右，这部分就是CUDA库和运行时的真实固定开销，不会出现上来就占4000+MiB的情况。

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