tensorflow限制GPU使用

c = np.random.random([5,8])#[vocab_size,embedding_size]
embedding = tf.Variable(c)
b = tf.nn.embedding_lookup(embedding, [1,4])#[batch_size,sentence_max_len](inputs)
print(b.shape)
#config = tf.ConfigProto()
#config.gpu_options.allow_growth = True
#config.gpu_options.per_process_gpu_memory_fraction = 0.1
#with tf.Session(config = config) as sess:
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print (sess.run(embedding))
    print (sess.run(b))

跑tensorflow的时候以为只是简单地测试tf.nn.embedding_lookup这个函数是怎么用的，并没有设置tf.ConfigProto(),以为不怎么占显存或占显存较少，但结果并不是这样的。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 387.34                 Driver Version: 387.34                    |
|-------------------------------+----------------------+----------------------+
| 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 |
|  0%   54C    P2    57W / 280W |  10743MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

记录可以看到显存基本被占满了，在notebook上restart kernel再查看显存

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 387.34                 Driver Version: 387.34                    |
|-------------------------------+----------------------+----------------------+
| 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 |
|  0%   50C    P8    17W / 280W |   2348MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

因此可以看到，为了避免过多地占用显存，应该根据使用需求限制GPU的资源使用,可以通过以下方法设置：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.1
with tf.Session(config = config) as sess:
    ...
    ...
    ...