tensorflow添加自定义OP（GPU版本）

最新推荐文章于 2025-06-06 00:14:05 发布

原创最新推荐文章于 2025-06-06 00:14:05 发布 · 3.4k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow自定义OP（GPU版本）

深度学习专栏收录该内容

3 篇文章

订阅专栏

今天试了下注册GPU支持的OP。
GPU内核
GPU内核分两部分实现：OpKernel和CUDA内核及其启动代码。

有时OpKernel的实现在CPU和GPU内核之间很常见，比如检查输入和分配输出。在这种情况下，建议的实施是：

定义在Device上模板化的OpKernel和张量的基本类型。
为了完成输出的实际计算，Compute函数调用模板函子结构。
该函数对CPUDevice的专门化定义在同一个文件中，但GPUDevice的专门化定义在.cu.cc文件中，因为它将与CUDA编译器一起编译。

在tensorflow/user_ops/下添加cuda_op_kernel.cu.cc，cuda_op_kernel.cc两个文件。

#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "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

这里定义的Eigen是一个高层次的C ++库，有效支持线性代数，矩阵和矢量运算，数值分析及其相关的算法。如果编译过程中说未定义eigen，可以自己在网上找个教程下载。然后将安装后产生的Eigen文件夹拷贝到/usr/local/include/下。

#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<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->template flat<int32>();

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

编译GPU设备的内核：
输入以下命令

1. TF_INC=$(python3.5 -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
2. TF_LIB=$(python3.5 -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())')
3. nvcc -std=c++11 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
4. -I $TF_INC -I$TF_INC/external/nsync/public -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

5. g++ -std=c++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
6. cuda_op_kernel.cu.o -I $TF_INC -fPIC -lcudart

输入第6步后可能会报： mutex.h:25:22: fatal error: nsync_cv.h: 没有那个文件或目录
在代码中添加nsync_cv.h的路径，修改后输入：

cuda_op_kernel.cu.o -I $TF_INC -I /usr/local/lib/python3.5/dist-packages/tensorflow/include/external/nsync/public/ -fPIC -lcudart

若报/usr/bin/ld: 找不到 -lcudart
那么你需要导入导入这个库文件的路径，如果未安装CUDA库，则/usr/local/lib64需要在上面的第6个命令中明确指定路径。例如，-L /usr/local/cuda-8.0/lib64/如果您的CUDA已安装，请添加/usr/local/cuda-8.0。

cuda_op_kernel.cu.o -I $TF_INC -I$TF_INC/external/nsync/public -fPIC -lcudart -L$TF_LIB -ltensorflow_framework -L/usr/local/cuda-8.0/lib64/cuda_op_kernel.cu.o -I $TF_INC -I /usr/local/lib/python3.5/dist-packages/tensorflow/include/external/nsync/public/ -fPIC -lcudart

生成cuda_op_kernel.so文件后你就可以使用tf.load_op_library(‘你的.so文件路径/cuda_op_kernel.so’)来使用你新添加的OP啦。

import tensorflow as tf

file = '/home/siat/Work/tensorflow-r1.4/tensorflow/user_ops/cuda_op_kernel.so'
cuda_op_module = tf.load_op_library(file)
with tf.Session(''):
  x = cuda_op_module.add_one([[6, 4], [2, 4]]).eval()
print(x)
#[[7 5]
 [3 5]]

注意：我用的Python版本是3.5，若你是使用ubuntu自带2.7版本的话请将命令行中python3.5改为python。

参考：
https://cloud.tencent.com/developer/section/1475696
https://blog.youkuaiyun.com/andylei777/article/details/78542624?locationNum=4&fps=1
https://blog.youkuaiyun.com/qq_27637315/article/details/79114633