tensorflowカスタムGPUバージョンopノード

この間、指導者は損失関数を修正するタスクを配置したが、この損失関数はtensorflowが持参したライブラリにはなく、多くの方法を考えてみたが、試してみると解決策が見つからなかった.tensorflowはフレームワークとデータを分離しているので、pythonで直接書いた関数は使えない.1つのノードを定義して呼び出すしかない.だから自然と1つのgpuバージョンのkernelルーチンを走ることを思い付いて、ネット上のcpuバージョンの教程は多くて、しかしgpuバージョンのは比較的に少なくて、公式サイトの教程の極の授業学院は話して、しかし私は話すのが複雑だと感じて、どうせ私は1回見て理解していないで、よし、本文を始めます.今回のルーチンでは,入力tensorの数値を加算して出力する.
手順1:kernelを書く
ファイル名:cuda_op_kernel.cu.ccコードは以下の通りです.

#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
__global__ void AddOneKernel(const int* in, const int N, int* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N;
       i += blockDim.x * gridDim.x) {
    out[i] = in[i] + 1;
  }
}
 
void AddOneKernelLauncher(const int* in, const int N, int* out) {
  AddOneKernel<<<32, 256>>>(in, N, out);
}
 
#endif

手順2:cppプログラムの作成
ファイル名:cuda_op_kernel.ccコードは以下の通りです.

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
 
using namespace tensorflow;
 
REGISTER_OP("AddOne")
    .Input("input: int32")
    .Output("output: int32")
    .Doc(R"doc(
Adds 1 to all elements of the tensor.
output: A Tensor.
  output = input + 1
)doc");
 
void AddOneKernelLauncher(const int* in, const int N, int* out);
 
class AddOneOp : public OpKernel {
 public:
  explicit AddOneOp(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();
 
    // 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->template flat();
 
    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    // Call the cuda kernel launcher
    AddOneKernelLauncher(input.data(), N, output.data());
  }
};
 
REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU), AddOneOp);

ステップ3:ノードのコンパイル
上の2つのフォルダの下で端末(ubuntu)を開き、次のコマンドを入力します.

TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')

g
++
-
std
=
c
++
11
-
shared
cuda_op_kernel
.cc
-
o
cuda_op_kernel
.so
-
fPIC
-
I
$TF_INC
-
O2
-
D_GLIBCXX_USE_CXX11_ABI
=
0
nvcc
-
std
=
c
++
11
-
c
-
o
cuda_op_kernel
.cu
.o
cuda_op_kernel
.cu
.cc
\
-
I
$TF_INC
-
D
GOOGLE_CUDA
=
1
-
x
cu
-
Xcompiler
-
fPIC
g
++
-
std
=
c
++
11
-
shared
-
o
cuda_op_kernel
.so
cuda_op_kernel
.cc
\
cuda_op_kernel
.cu
.o
-
I
$TF_INC
-
fPIC
-
lcudart
4番目のコマンドを実行すると、次のようなエラーが発生する可能性があります.

/usr/bin/ld:     -lcudart
collect2: error: ld returned 1 exit status

この問題について私は一日やったが、この言葉の意味はlibcudartというものがあることを明らかにした.soのライブラリファイルが見つからないので、/usr/bin/ディレクトリの下に行ってlibcudartという名前のファイルが見つかりませんでした.soのライブラリファイル、それではどうすればいいですか.じゃ、パソコンの中でこのライブラリファイルを探します.私のパソコンのライブラリファイルはこのフォルダの下にあります./usr/local/cuda/targets/x 86_64-linux/lib、cudaというファイルの下で検索してもいいです.私はこのフォルダの下にあると思います.住所を知ってから、元の4番目のコマンドを次のコマンドに置き換えます.

g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc cuda_op_kernel.cu.o -L/usr/local/cuda/targets/x86_64-linux/lib -I $TF_INC -fPIC -lcudart

はコンパイルに合格しました.
テストファイルを書きました
手順4:テスト
ファイル名:test.pyコードは次のとおりです.

import tensorflow as tf
cuda_op_module = tf.load_op_library('./cuda_op_kernel.so')
with tf.Session(''):
  x=cuda_op_module.add_one([[1, 2], [3, 4]]).eval()

print x

出力Yes:
[[2 3]
 [4 5]]