本来想写一篇在windows下c/c++版tensorflow在多GPU设备下进行并行推理的文章的,但由于之前配置测试的时候没有随手进行截图和记录的习惯,导致严重缺少素材,现在再进行复现就很难了,哎,我要改掉这个毛病。下面进入正题。
TensorFlow 在运行中,通常默认占用机器上的所有 GPU 资源,但实际运行模型并不需要占用如此多的资源,为了使 GPU 资源得到充分的使用,我们通常需要手动指定 TensorFlow 使用的 GPU 资源,在使用 Python 进行 TensorFlow 开发时,对于 GPU 资源的设置很方便,但是在使用 C/C++ 对 TensorFlow 模型进行工程化实现时,这个设置就变得相对复杂。由于近期在对 TensorFlow 进行 C/C++ 工程化时遇到了指定 GPU 资源的问题,于是总结一下。
c_api 设置 GPU 需要经过两个步骤:
1.在 Python 端设置 GPU 选项,然后把配置转化为序列化的 protobuf 字符串;
2.将序列化的 protobuf 字符串提供给 c_api 中的 TF_SetConfig() 函数。
下面举例说明:
Python 端设置
tensorflow 中 GPU 的选项描述,可以到这里查看,具体这里就不详细描述。对于 visible_device_list 选项的描述,具体在这里有描述。
python代码:
比如设置显存占用比例为0.333,允许根据需要自动扩展内存,可见GPU设备为"0".
python
import tensorflow as tf
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333, allow_growth=True, visible_device_list='0')
config =tf.ConfigProto(gpu_options=gpu_options)
serialized = config.SerializeToString()
list(map(hex, serialized))
输出结果为:
['0x32', '0xe', '0x9', '0x1d', '0x5a', '0x64', '0x3b', '0xdf', '0x4f', '0xd5', '0x3f', '0x20', '0x1', '0x2a', '0x1', '0x30']
c_api 代码中设置
将上面得到的指令赋值给config,然后通过TF_SetConfig去对GPU进行配置。
cpp
TF_Graph* graph = TF_NewGraph ();
TF_Status* status = TF_NewStatus ();
TF_SessionOptions*session_opts = TF_NewSessionOptions();
uint8_t config[16] ={0x32, 0xe, 0x9, 0x1d, 0x5a, 0x64, 0x3b, 0xdf, 0x4f, 0xd5, 0x3f, 0x20, 0x1, 0x2a, 0x1, 0x30};
TF_SetConfig(session_opts, (void*)config, 16, status);
TF_Session* sess = TF_NewSession(graph, session_opts, status);
TensorFlow C++ API 指定 GPU
c++版本tensorflow中就比c版方便多了,直接通过api去进行设置。
cpp
tensorflow::SessionOptions session_options;
tensorflow::graph::SetDefaultDevice("/gpu:0", &graph_def);
session_options.config.mutable_gpu_options()->set_allow_growth(true);
session_options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.5);
另外提供C版tensorflow的下载地址:
安装C版tensorflow