如何监控TensorFlow Serving的GPU内存？查询模型实际占用及剩余可用内存

好问题！这是在生产环境部署TensorFlow Serving时非常常见的需求，我来一步步帮你解决这些问题：

你目前用nvidia-smi查看进程级内存是最直观的方式，但还有更细粒度、更适合生产环境的方法：

• 启用TensorFlow Serving的监控指标：启动TF Serving时，加上--enable_batching --monitoring_config_file=monitoring_config.pbtxt参数（需要提前编写一个简单的配置文件），这样服务会暴露Prometheus格式的监控指标，其中包含tensorflow_serving_model_analysis_gpu_memory_usage这类直接关联GPU内存使用的数据，方便你集成到监控系统中实时追踪。
• nvidia-smi的进阶用法：如果不想搞复杂的监控配置，也可以用nvidia-smi pmon -s u实时监控进程的GPU内存使用率，或者nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv输出结构化的进程内存数据，方便脚本自动解析和告警。
• 内置内存分析工具：如果你的TF Serving是自定义镜像，可以在启动脚本中加入TensorFlow的内存追踪代码，比如tf.config.experimental.get_memory_info('GPU:0')，能拿到进程内部更详细的显存分配细节，包括已分配、已使用的显存数值。

nvidia-smi只能看到整个TF Serving进程的总显存占用，要拿到单个模型的实际内存，得用这些方法：

• 查看模型加载日志：启动TF Serving时确保日志级别为INFO（默认就是），并设置TF_CPP_MIN_LOG_LEVEL=0让TensorFlow输出详细日志。当模型加载完成时，日志里会出现类似Total memory usage of variables: X MiB的条目，这就是该模型参数的内存占用（注意：这只是静态参数的内存，不包括推理时产生的临时张量显存）。
• 自定义内存差值计算：如果需要更精确的单个模型内存（包括初始化时的额外显存），可以在自定义TF Serving的加载逻辑中，在加载模型前后分别调用tf.config.experimental.get_memory_info('GPU:0')，两者的差值就是这个模型实际占用的显存。
• TF Lite模型专属方法：如果你的模型是TF Lite格式，可以用tf.lite.experimental.get_memory_usage直接输出模型在GPU上的内存占用，但这个只适用于TF Lite场景，TF Serving默认用的是SavedModel。

要确定能不能安全加载新模型，需要结合当前空闲显存和新模型的预估内存来判断：

• 先预估新模型的内存需求：在本地测试环境加载该模型，通过日志或者tf.config.experimental.get_memory_info拿到模型参数的内存，再额外预留20%-30%的显存作为推理时的临时空间（推理过程中会产生中间张量，也需要显存）。
• 获取当前GPU空闲显存：用nvidia-smi --query-gpu=free_memory --format=csv,noheader,nounits可以直接获取空闲显存的数值；如果配置了NVIDIA的Prometheus exporter，也可以通过nvidia_gpu_memory_free指标来获取。
• 自动化判断逻辑：可以写一个简单的脚本，定期拉取当前空闲显存，和新模型的预估内存（加上预留空间）对比，如果空闲显存大于等于这个值，就可以安全加载。另外，TF Serving加载模型时如果显存不足，会抛出ResourceExhaustedError，也可以尝试加载并捕获这个异常，但这种是事后判断，不如提前预估可靠。
• 开启动态显存分配：启动TF Serving时设置环境变量TF_FORCE_GPU_ALLOW_GROWTH=true，这样TensorFlow会根据实际需求动态分配显存，而不是一开始就占满整个GPU，能更高效利用显存，也降低了加载新模型时显存不足的概率。

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