如何在 Tensorflow C++ 中设置 inter_op_parallelism_threads 和 intra_op_parallelism_threads 以进行 Seesion 配置

Posted 2023-02-23

【中文标题】如何在 Tensorflow C++ 中设置 inter_op_parallelism_threads 和 intra_op_parallelism_threads 以进行 Seesion 配置

【英文标题】：how to set inter_op_parallelism_threads and intra_op_parallelism_threads for seesion configuration in Tensorflow c++

【发布时间】：2020-04-22 09:10:08

【问题描述】：

我知道如何在python tensorflow中设置内部线程和内部线程，如下所示

n_cpus=1
sess = tf.Session(config=tf.ConfigProto(
     device_count= "CPU": n_cpus ,
    inter_op_parallelism_threads=n_cpus,
    intra_op_parallelism_threads=8
));

但我需要在 C++ 中实现它。所以在 tensorflow c++ 中设置线程间和内部线程的任何等效代码？？

【问题讨论】：

您可以查看this，评论中有示例代码，希望对您有所帮助：D

【参考方案1】：

您可以使用以下代码在 C++ 中实现这一点。

Status LoadGraph(string graph_file_name, std::unique_ptr<tensorflow::Session>* session) 
  tensorflow::GraphDef graph_def;
  Status load_graph_status =
      ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
  if (!load_graph_status.ok()) 
    return tensorflow::errors::NotFound("Failed to load compute graph at '",
                                        graph_file_name, "'");
  
  
  // initialize the number of worker threads
    tensorflow::SessionOptions options;
    tensorflow::ConfigProto & config = options.config;
    config.set_inter_op_parallelism_threads(1);
    config.set_intra_op_parallelism_threads(1);
    config.set_use_per_session_threads(false); 

  session->reset(tensorflow::NewSession(options));
  Status session_create_status = (*session)->Create(graph_def);
  if (!session_create_status.ok()) 
    return session_create_status;
  

  return Status::OK();