笔者需要tensorflow仅运行在一个GPU上(机器本身有多GPU),而且需要依据系统参数动态调节,故无法简单使用CUDA_VISIBLE_DEVICES
。
一种方式是全局使用tf.device
函数生成的域,但设备号需要在绘制Graph前指定,仍然不够灵活。
查阅文档发现config
的GPUOptions
中的visible_device_list
可以定义GPU编号从visible到virtual的映射,即可以设置tensorflow可见的GPU device,从而全局设置了tensorflow可见的GPU编号。代码如下:
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(device_num)
sess = tf.Session(config=config)
参考 多卡服务器下隐藏部分 GPU 和 TensorFlow 的显存使用设置,还可以通过os包设置全局变量CUDA_VISIBLE_DEVICES
,代码如下:
os.environ["CUDA_VISIBLE_DEVICES"] = "2"