tensorflow指定GPU与动态分配GPU memory设置

发布时间:2020-02-06 08:42 来源:互联网 当前栏目:网站服务器

在tensorflow中,默认指定占用所有的GPU,如需指定占用的GPU,可以在命令行中:

export CUDA_VISIBLE_DEVICES=1

这样便是只占用1号GPU,通过命令

nvidia-smi

可以查看各个GPU的使用情况。

另外,也可以在python程序中指定GPU,并且动态分配memory,代码如下

import os
import sys
os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
set_session(tf.Session(config=config))

这样,占用的GPU就是sys.argv[1]了,并且会按需分配memory。

以上这篇tensorflow指定GPU与动态分配GPU memory设置就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持IIS7站长之家。

  • 1、
  • 2、
  • 3、
  • 4、
  • 5、
  • 6、
  • 7、
  • 8、
  • 9、
  • 10、
  • 11、
  • 12、
  • 13、
  • 14、
  • 15、
  • 16、
  • 17、
  • 18、
  • 19、
  • 20、
  • 21、
  • 22、
  • 23、
  • 24、
  • 25、
  • 1、
  • 2、
  • 3、
  • 4、
  • 5、
  • 6、
  • 7、
  • 8、
  • 9、
  • 10、
  • 11、
  • 12、
  • 13、
  • 14、
  • 15、
  • 16、
  • 17、
  • 18、
  • 19、
  • 20、
  • 21、
  • 22、
  • 23、
  • 24、
  • 25、