如何测量(手动)我的GPU内存已使用/可用

Question

我正在通过创建不同的模型并对其进行测试来尝试Tensorflow。我现在遇到的问题是，在我面对OOM (内存不足)错误之前，我不清楚我的模型可能有多大。当然，我可以继续添加图层或者让它们变得更大，看看什么时候我达到了极限。但如果我能测量我的模型占用了多少内存，那就太好了。

使用nvidia-smi命令，我可以看到：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 430.40       Driver Version: 430.40       CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  Off  | 00000000:01:00.0 Off |                  N/A |
| 14%   56C    P2    87W / 280W |  10767MiB / 11178MiB |     37%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     24329      C   python                                     10747MiB |
+-----------------------------------------------------------------------------+

但是我这里的内存使用的问题是，无论我的模型是大是小，它总是显示相同的数字！似乎Tensorflow无论如何都会得到一切。

我正在寻找一个公式，它通过提供模型参数的数量、批量大小和样本大小来为我提供所需的内存。

Answer 1

您可以使用next函数：

def get_max_memory_usage(sess):
    """Might be unprecise. Run after training"""
    if sess is None: sess = tf.get_default_session()
    max_mem = int(sess.run(tf.contrib.memory_stats.MaxBytesInUse()))
    print("Contrib max memory: %f G" % (max_mem / 1024. / 1024. / 1024.))
    return max_mem

或者看看弗拉迪斯拉夫·布拉托夫的“mem_utils”