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

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

本文关键字：tensorflow GPU op 自定义 运行 C++ 代码      更新时间：2023-10-16

我按照这些例子在TensorFlow中编写自定义操作：
添加新的操作
cuda_op_kernel将
函数更改为我需要执行的操作
，但所有示例都是Python代码中的测试。
我需要从 c++ 代码运行我的操作，我该怎么做？

这个简单的示例显示了使用 API 构建和执行图形C++：

// 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，则必须使用几乎相同的代码。我在这个存储库中有一个自定义操作，我将用作示例代码。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 文件中的 Cuda 内核。要启动 Op，我必须(再次(创建一个新的计算图，注册我的操作，打开一个会话并使其从我的代码运行：

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>();
相关文章：