作者:丁志翔64164 | 来源:互联网 | 2023-02-01 20:39
我正在尝试在本文档之后宽松地向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();
// 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);
cuda_op_kernel.cu
#define EIGEN_USE_GPU
#include
#include
__global__ void AddOneKernel(const int* in, const int N, int* out) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i >>(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)
test.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 和DEVICE_GPU
与DEVICE_CPU
,运算输出右值[6 5 4 3 2]
(具有完全一样的CMakeList.txt
)。
任何想法如何获取正确的值以返回?