如何从 C++ 代码运行自定义 GPU tensorflow::op? [英] How to run custom GPU tensorflow::op from C++ code?

问题描述

我按照以下示例在 TensorFlow 中编写自定义操作:
添加新操作
cuda_op_kernel
把功能改成我需要做的操作.
但所有示例都是 Python 代码中的测试.
我需要从 C++ 代码运行我的操作，我该怎么做?

I follow these examples to write custom op in TensorFlow:
Adding a New Op
cuda_op_kernel
Change the function to operation I need to do.
But all the examples are tests in Python code.
I need to run the my op from c++ code, how can I do this?

推荐答案

这个简单的例子展示了使用 C++ API:

This simple example shows the construction and the execution of a graph using C++ API:

// tensorflow/cc/example/example.cc

#include "tensorflow/cc/client/client_session.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor.h"

int main() {
  using namespace tensorflow;
  using namespace tensorflow::ops;
  Scope root = Scope::NewRootScope();
  // Matrix A = [3 2; -1 0]
  auto A = Const(root, { {3.f, 2.f}, {-1.f, 0.f} });
  // Vector b = [3 5]
  auto b = Const(root, { {3.f, 5.f} });
  // v = Ab^T
  auto v = MatMul(root.WithOpName("v"), A, b, MatMul::TransposeB(true)); // <- in your case you should put here your custom Op
  std::vector<Tensor> outputs;
  ClientSession session(root);
  // Run and fetch v
  TF_CHECK_OK(session.Run({v}, &outputs));
  // Expect outputs[0] == [19; -3]
  LOG(INFO) << outputs[0].matrix<float>();
  return 0;
}

就像在 Python 中一样，您首先需要在一个范围内构建一个计算图，在这种情况下，它只有一个矩阵乘法，其终点在 v 中.然后您需要为范围打开一个新会话 (session)，并在您的图形上运行它.在这种情况下，没有提要字典，但在页面末尾有一个关于如何提要值的示例:

As in the Python counterpart, you first need to build a computational graph in a scope, which in this case has only a matrix multiplication in it, whose end point is in v. Then you need to open a new session (session) for the scope, and run it on your graph. In this case there is no feed dictionary, but at the end of the page there is an example on how to feed values:

Scope root = Scope::NewRootScope();
auto a = Placeholder(root, DT_INT32);
// [3 3; 3 3]
auto b = Const(root, 3, {2, 2});
auto c = Add(root, a, b);
ClientSession session(root);
std::vector<Tensor> outputs;

// Feed a <- [1 2; 3 4]
session.Run({ {a, { {1, 2}, {3, 4} } } }, {c}, &outputs);
// outputs[0] == [4 5; 6 7]

^{此处报告的所有代码段均来自 TensorFlow 的 C++ API 指南}

如果您想调用自定义 OP，您必须使用几乎相同的代码.我在 这个存储库 中有一个自定义操作，我将用作示例代码.OP 已注册:

If you want to call custom OP you have to use almost the same code. I have a custom op in this repository that I will use as an example code. The OP has been registered:

REGISTER_OP("ZeroOut")
  .Input("to_zero: int32")
  .Output("zeroed: int32")
  .SetShapeFn([](::tensorflow::shape_inference::InferenceContext *c) {
    c->set_output(0, c->input(0));
    return Status::OK();
  });

并且 Op 被定义为 cuda 文件.要启动 Op，我必须(再次)创建一个新的计算图，注册我的 op，打开一个会话并让它从我的代码中运行:

and the Op is defined to be a Cuda Kernel in the cuda file. To launch the Op I have to (again), create a new computational graph, register my op, open a session and make it run from my code:

Scope root = Scope::NewRootScope();
// Matrix A = [3 2; -1 0]
auto A = Const(root, { {3.f, 2.f}, {-1.f, 0.f} });
auto v = ZeroOut(root.WithOpName("v"), A); 
std::vector<Tensor> outputs;
ClientSession session(root);
// Run and fetch v
TF_CHECK_OK(session.Run({v}, &outputs));
LOG(INFO) << outputs[0].matrix<float>();

这篇关于如何从 C++ 代码运行自定义 GPU tensorflow::op?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！