Tensorflow 如何获得Keras的确切GPU内存使用率

我最近开始学习Keras和TensorFlow。我正在测试MNIST数据集上的一些模型（非常基本的东西）。我想知道，在训练和推理过程中，我的模型到底消耗了多少内存。我试着在谷歌上搜索，但没有找到多少信息

我遇到了Nvidia smi。我尝试使用config.gpu_options.allow_growth=True选项，但由于Nvidia smi的一些问题，仍然无法使用python.exe消耗的确切内存。我知道我可以单独运行火车和推理，但这太麻烦了。如果我能找到合适的API来完成这项工作，那就很容易了

Tensorflow是一个如此著名和使用广泛的库，我希望找到一种更好更快的方法来获取这些数字

最后，我的问题是： 如何在训练和推理过程中获得Keras模型的准确内存使用情况

相关规格： 操作系统：Windows 10 GPU:GTX 1050 TensorFlow版本：1.14

如果需要任何其他细节，请告诉我

---

谢谢

在TF1中有函数

tf.contrib.memory\u stats.MaxBytesInUse（）

。不确定它是否仍包含在TF2中。谢谢，我将对此进行查看。@Raviteja Narra，您能否确认它是否已通过上述注释得到解决？