Tensorflow C++ GPU设置显存

Tensorflow C++ GPU设置显存

Tensorflow 在运行中经常会占用全部的显存，并且在C++工程化时会遇到不自动释放显存，本文将介绍如何在C++中设置GPU显存。

Python端设置显存的方式

• 在构造tf.Session()时候通过传递tf.GPUOptions作为可选配置参数的一部分来显式地指定需要分配的显存比例，如下所示：在构造tf.Session()时候通过传递tf.GPUOptions作为可选配置参数的一部分来显式地指定需要分配的显存比例，如下所示：

  import tensorflow as tf
  gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
  sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

• per_process_gpu_memory_fraction指定了每个GPU进程中使用显存的上限，但它只能均匀作用于所有GPU，无法对不同GPU设置不同的上限.

C++ 端设置

• 对于TF_SetConfig中，需要提供把配置作为序列化的protobuf字符串处理。在python端获取字符串

  serialized = config.SerializeToString()
  list(map(hex, serialized))

• 序列化结果:

  [‘0x32’, ‘0x5’, ‘0x20’, ‘0x1’, ‘0x2a’, ‘0x1’, ‘0x33’]

• C++设置：

  TF_Graph* graph = TF_NewGraph ();
  TF_Status* status = TF_NewStatus ();
  TF_SessionOptionssession_opts = TF_NewSessionOptions();
  uint8_t config[7] ={0x32, 0x5, 0x20, 0x1, 0x2a, 0x1, 0x32};
  TF_SetConfig(session_opts,(void)config,7,status);
  TF_Session* sess = TF_NewSession(graph, session_opts, status);