tensor flow add op

- op interface

인풋, 아웃풋의 타입과 이름을 정해준다. 배열 또는 행렬의 크기는 동적으로 할당 된다.

op에 따라서는 docstring과 attrs도 정해준다.

tensorflow/core/user_ops/new_op.cc을 만들어준다.

#include "tensorflow/core/framework/op.h"

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

ZeroOut이라는 op를 등록했다. 

인풋의 이름은 to_zero이며 int32을 입력 받는다.

아웃풋의 이름은 zeroed이며 int32을 출력 한다.

op의 이름은 CamelCase이며 유일해야한다. 그리고 _로 시작하면 안된다 내부적으로 이미사용하고있다.

- kernel for the op

OpKernel을 상속받은 클래스를 만들어야한다.

그리고 Compute을 override해야한다.

Compute함수는 하나의 OpKernelContext* context를 인자로 받는데 

이 것을 통해서 입력과 출력에 접근 할 수 있다.

Compute함수는 쓰레드환경에서 동작하기때문에 유념하며 작성해야한다.

mutex을 사용하거나 클래스 변수를 공유하지 말아라.

ResourceMgr을 이용해서 op의 상태를 확인 할 수 있다.

include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32");
	
class ZeroOutOp : public OpKernel 
{
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override 
  {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) 
    {
      output(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output(0) = input(0);
  }
};
REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

 코드를 보면 context을 받아 input 텐서를 뽑고 flat한다. 입력에 기반한 출력의 크기를 OP_REQUIRES_OK을 이용하여

output_tensor에 할당했다. 그리고 output_tensor을 flat하여 첫번째 인자만 출력으로 전달하고 나머지는 0으로 한다.

kernel을 만들었으면 텐서플로우에 등록을 해야한다. 

REGISTER_KERNEL_BUILDER을 이용하여 CPU에서 ZeroOut이라는 op는 ZeroOutOp클래스로 연결된다.

- Building the op library with tensorflow source installation

BUILD라는 파일을 만든다. 그리고 아래 처럼 안에 쓴다.

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

zero_out.cc을 zero_out.so로 빌드하는 것을 알 수 있다. 아래 명령어를 터미널에 치면 빌드 된다.

$ bazel build -c opt //tensorflow/core/user_ops:zero_out.so

- Using the Op in Python

텐서플로우 파이썬