5

我正在尝试在文档之后松散地向 TensorFlow 添加一个新操作。不同之处在于我正在尝试实现基于 GPU 的操作。我要添加的操作是来自此处的 cuda 操作(cuda_op.py、cuda_op_kernel.cc、cuda_op_kernel.cu.cc)。我正在尝试在 tensorflow 之外编译这些并tf.load_op_library用于将它们拉入。我做了一些更改,所以这里是我的文件:

cuda_op_kernel.cc

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

using namespace tensorflow;  // NOLINT(build/namespaces)

REGISTER_OP("AddOne")
    .Input("input: int32")
    .Output("output: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

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

cuda_op_kernel.cu

#define EIGEN_USE_GPU
#include <cuda.h>
#include <stdio.h>

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

  cudaError_t cudaerr = cudaDeviceSynchronize();
  if (cudaerr != cudaSuccess)
    printf("kernel launch failed with error \"%s\".\n", cudaGetErrorString(cudaerr));
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.5)

#found from running python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())'
include_directories(/usr/local/lib/python3.5/dist-packages/tensorflow/include)

find_package(CUDA)

#set flags based on tutorial
set (CMAKE_CXX_FLAGS "--std=c++11 -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0")

#pass flags to c++ compiler
SET(CUDA_PROPAGATE_HOST_FLAGS ON)

#create library
cuda_add_library(
    cuda_op SHARED
    src/cuda_op_kernel.cu
    src/cuda_op_kernel.cc
    OPTIONS -gencode=arch=compute_20,code=sm_20)

#copy test file to build folder
configure_file(src/test.py test.py COPYONLY)

测试.py

import tensorflow as tf
mod = tf.load_op_library('./libcuda_op.so')
with tf.Session() as sess:
    start = [5,4,3,2,1]
    print(start)
    print(mod.add_one(start).eval())

我能够test.py成功编译和运行,但输出始终是[0 0 0 0 0]. 如果我AddOneKernel<<<32, 256>>>(in, N, out);for (int i = 0; i < N; i++) out[i] = in[i] + 1;DEVICE_GPU替换DEVICE_CPU,则运算输出正确的值[6 5 4 3 2](完全相同CMakeList.txt)。

知道如何获得要返回的正确值吗?

4

3 回答 3

3

我不完全记得我在哪里找到了 CUDA 的 cmake 东西,但是这些选项以某种方式弄乱了编译。替换cuda_add_libraryCMakeLists.txt以下解决了该问题。

#no options needed
cuda_add_library(
    cuda_op SHARED
    src/cuda_op_kernel.cu
    src/cuda_op_kernel.cc)
于 2017-06-07T16:56:15.590 回答
1

在Tensorflow上查看当前官方 GPU 操作构建说明,添加 GPU 操作支持

nvcc -std=c++11 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -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 ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

正如它所说,请注意,如果您的 CUDA 库未安装在 中,则/usr/local/lib64需要在上面的第二个 (g++) 命令中明确指定路径。例如,-L /usr/local/cuda-8.0/lib64/如果您的 CUDA 安装在/usr/local/cuda-8.0.

此外,请注意在某些 linux 设置中,需要 nvcc 编译步骤的附加选项。添加-D_MWAITXINTRIN_H_INCLUDED到 nvcc 命令行以避免来自mwaitxintrin.h.

于 2018-09-27T03:56:25.860 回答
0

ubuntu@cubuntu:~/Desktop/src/src/build$ cmake ..

-- 配置完成

-- 生成完成

-- 构建文件已写入:/home/ubuntu/Desktop/src/src/build

ubuntu@cubuntu:~/Desktop/src/src/build$ make

[ 33%] 构建 NVCC(设备)对象 CMakeFiles/cuda_op.d/cuda_op_generated_cuda_op_kernel.cu.o

nvcc 警告:“compute_20”、“sm_20”和“sm_21”架构已弃用,可能会在未来版本中删除(使用 -Wno-deprecated-gpu-targets 来抑制警告)。

nvcc 警告:“compute_20”、“sm_20”和“sm_21”架构已弃用,可能会在未来版本中删除(使用 -Wno-deprecated-gpu-targets 来抑制警告)。

扫描目标 cuda_op 的依赖关系

[ 66%] 构建 CXX 对象 CMakeFiles/cuda_op.dir/cuda_op_kernel.cc.o /home/ubuntu/Desktop/src/src/cuda_op_kernel.cc:1:17:错误:“tensorflow”不是使用命名空间的命名空间名称张量流;// NOLINT(构建/命名空间)

于 2017-09-26T16:05:30.420 回答