CODEWITHSUNDEEP
c++tensorflow
alon
alon

Reputation: 240

How to run custom GPU tensorflow::op from C++ code?

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?

Upvotes: 2

Views: 2394

Answers (1)

Matteo Ragni
Matteo Ragni

Reputation: 2956

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;
}

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]

All the code segments here reported come from the C++ API guide for TensorFlow

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

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

Upvotes: 2

Related Questions