使用tensorflow,如果不加設置,即使是很小的模型也會占用整塊GPU,造成資源浪費。
所以我們需要設置,使程序按需使用GPU。
具體設置方法:
1 gpu_options = tf.GPUOptions(allow_growth=True) 2 sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) 3 # sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))
1 gpu_no = '0' # or '1' 2 os.environ["CUDA_VISIBLE_DEVICES"] = gpu_no 3 4 # 定義TensorFlow配置 5 config = tf.ConfigProto() 6 7 # 配置GPU內存分配方式,按需增長,很關鍵 8 config.gpu_options.allow_growth = True 9 10 # 配置可使用的顯存比例 11 config.gpu_options.per_process_gpu_memory_fraction = 0.1 12 13 # 在創建session的時候把config作為參數傳進去 14 sess = tf.InteractiveSession(config = config)
1 gpu_options = tf.GPUOptions(allow_growth=True) 2 # sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) 3 4 with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
說明:使用jupyter notebook,如果沒有上述設置,整個GPU會一直被占用;使用pycharm,雖然GPU不會被一直占用,但是運行時還是會占用整個GPU,所以也需要上述設置。