如何从 C++ 代码运行自定义 GPU tensorflow::op？

Question

我按照这些示例在 TensorFlow 中编写自定义操作：
Adding a New Op
cuda_op_kernel
将功能改为我需要做的操作。
但所有示例都是 Python 代码中的测试。
我需要从 c++ 代码运行我的操作，我该怎么做？

Answer 1

这个简单的例子展示了使用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)，并在您的图表上运行它。在这种情况下，没有提要字典，但在页面末尾有一个关于如何提要值的示例：

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，你必须使用几乎相同的代码。我在this repository 中有一个自定义操作，我将使用它作为示例代码。 OP 已注册：

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 file 中的 Cuda 内核。要启动操作，我必须（再次）创建一个新的计算图，注册我的操作，打开一个会话并让它从我的代码中运行：

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>();