在tensorflow中,默認指定占用所有的GPU,如需指定占用的GPU,可以在命令行中:
export CUDA_VISIBLE_DEVICES=1
這樣便是只占用1號GPU,通過命令
nvidia-smi
可以查看各個GPU的使用情況。
另外,也可以在python程序中指定GPU,并且動態分配memory,代碼如下
import osimport sysos.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]import tensorflow as tffrom keras.backend.tensorflow_backend import set_sessionconfig = tf.ConfigProto()config.gpu_options.allow_growth = Trueset_session(tf.Session(config=config))
這樣,占用的GPU就是sys.argv[1]了,并且會按需分配memory。
以上這篇tensorflow指定GPU與動態分配GPU memory設置就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支持武林站長站。
新聞熱點
疑難解答