TensorFlow C/C++ API设置GPU选项

TensorFlow在运行中,通常默认占用机器上的所有GPU资源,但实际运行模型并不需要占用如此多的资源,为了使GPU资源得到充分的使用,我们通常需要手动指定TensorFlow使用的GPU资源,在使用Python进行TensorFlow开发时,对于GPU资源的设置很方便,但是在使用C/C++对TensorFlow模型进行工程化实现时,这个设置就变得相对复杂。由于近期在对TensorFlow进行C/C++工程化时遇到了指定GPU资源的问题,于是总结一下。感谢网友lucas的帮助

TensorFlow c_api 指定GPU资源

c_api设置GPU需要经过两个步骤:

    1. 在Python端设置GPU选项,然后把配置转化为序列化的protobuf字符串;
    1. 将序列化的protobuf字符串提供给c_api中的TF_SetConfig()函数。

下面举例说明:

  • Python端设置
    tensorflow中GPU的选项描述,可以到这里查看,具体这里就不详细描述。对于visible_device_list选项的描述,具体在这里有描述。
1
2
3
4
5
6
7
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))

输出结果为

1
['0x32', '0xe', '0x9', '0x1d', '0x5a', '0x64', '0x3b', '0xdf', '0x4f', '0xd5', '0x3f', '0x20', '0x1', '0x2a', '0x1', '0x30']

  • c_api代码中设置
1
2
3
4
5
6
7
8
9
TF_Graph* graph = TF_NewGraph ();
TF_Status* status = TF_NewStatus ();
TF_SessionOptions*session_opts = TF_NewSessionOptions();
/*PS:此处要将序列化的protobuf结果中的单引号去掉,直接设置16进制数*/
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

1
2
3
4
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);

参考文献

-------------本文结束感谢您的阅读-------------
分享