title: how to programmatically determine available GPU memory with tensorflow?
body: <p>For a vector quantization (k-means) program I like to know the amount of available memory on the present GPU (if there is one). This is needed to choose an optimal batch size in order to have as few batches as possible to run over the complete data set.</p>
  
  <p>I have written the following test program:</p>
  
  <pre><code>import tensorflow as tf
  import numpy as np
  from kmeanstf import KMeansTF
  print("GPU Available: ", tf.test.is_gpu_available())
  
  nn=1000
  dd=250000
  print("{:,d} bytes".format(nn*dd*4))
  dic = {}
  for x in "ABCD":
      dic[x]=tf.random.normal((nn,dd))
      print(x,dic[x][:1,:2])
  
  print("done...")
  </code></pre>
  
  <p>This is a typical output on my system with (ubuntu 18.04 LTS, GTX-1060 6GB). Please note the core dump.</p>
  
  <pre><code>python misc/maxmem.py 
  GPU Available:  True
  1,000,000,000 bytes
  A tf.Tensor([[-0.23787294 -2.0841186 ]], shape=(1, 2), dtype=float32)
  B tf.Tensor([[ 0.23762687 -1.1229591 ]], shape=(1, 2), dtype=float32)
  C tf.Tensor([[-1.2672468   0.92139906]], shape=(1, 2), dtype=float32)
  2020-01-02 17:35:05.988473: W tensorflow/core/common_runtime/bfc_allocator.cc:419] Allocator (GPU_0_bfc) ran out of memory trying to allocate 953.67MiB (rounded to 1000000000).  Current allocation summary follows.
  2020-01-02 17:35:05.988752: W tensorflow/core/common_runtime/bfc_allocator.cc:424] **************************************************************************************************xx
  2020-01-02 17:35:05.988835: W tensorflow/core/framework/op_kernel.cc:1622] OP_REQUIRES failed at cwise_ops_common.cc:82 : Resource exhausted: OOM when allocating tensor with shape[1000,250000] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
  Segmentation fault (core dumped)
  </code></pre>
  
  <p>Occasionally I do get an error from python instead of a core dump (see below). This would actually be better since I could catch it and thus determine by trial and error the maximum available memory. But it alternates with core dumps:</p>
  
  <pre><code>python misc/maxmem.py 
  GPU Available:  True
  1,000,000,000 bytes
  A tf.Tensor([[-0.73510283 -0.94611156]], shape=(1, 2), dtype=float32)
  B tf.Tensor([[-0.8458411  0.552555 ]], shape=(1, 2), dtype=float32)
  C tf.Tensor([[0.30532074 0.266423  ]], shape=(1, 2), dtype=float32)
  2020-01-02 17:35:26.401156: W tensorflow/core/common_runtime/bfc_allocator.cc:419] Allocator (GPU_0_bfc) ran out of memory trying to allocate 953.67MiB (rounded to 1000000000).  Current allocation summary follows.
  2020-01-02 17:35:26.401486: W tensorflow/core/common_runtime/bfc_allocator.cc:424] **************************************************************************************************xx
  2020-01-02 17:35:26.401571: W tensorflow/core/framework/op_kernel.cc:1622] OP_REQUIRES failed at cwise_ops_common.cc:82 : Resource exhausted: OOM when allocating tensor with shape[1000,250000] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
  Traceback (most recent call last):
    File "misc/maxmem.py", line 11, in &lt;module&gt;
      dic[x]=tf.random.normal((nn,dd))
    File "/home/fritzke/miniconda2/envs/tf20b/lib/python3.7/site-packages/tensorflow_core/python/ops/random_ops.py", line 76, in random_normal
      value = math_ops.add(mul, mean_tensor, name=name)
    File "/home/fritzke/miniconda2/envs/tf20b/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py", line 391, in add
      _six.raise_from(_core._status_to_exception(e.code, message), None)
    File "&lt;string&gt;", line 3, in raise_from
  tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[1000,250000] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:Add] name: random_normal/
  </code></pre>
  
  <p>How could I reliably get this information for whatever system the software is running on?</p>