[英]Adding a GPU Op in Tensorflow
我正在嘗試按照此文檔松散地向TensorFlow添加新操作。 不同之處在於我正在嘗試實現基於GPU的操作。 我試圖添加的操作是來自這里的cuda op(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)
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);
with for (int i = 0; i < N; i++) out[i] = in[i] + 1;
和DEVICE_GPU
一起使用DEVICE_CPU
,op輸出正確的值[6 5 4 3 2]
(具有完全相同的CMakeList.txt
)。
知道如何獲得正確的值返回?
我不完全記得我在哪里找到了CUDA的cmake東西,但選項卻以某種方式弄亂了編譯。 將CMakeLists.txt
cuda_add_library
替換為以下修復了該問題。
#no options needed
cuda_add_library(
cuda_op SHARED
src/cuda_op_kernel.cu
src/cuda_op_kernel.cc)
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'不是使用命名空間的命名空間名稱tensorflow; // NOLINT(構建/命名空間)
查看Tensorflow上當前正式的GPU ops構建說明, 添加GPU op支持
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 ++)命令中明確指定路徑。 例如,如果您的CUDA安裝在/usr/local/cuda-8.0
,請添加-L /usr/local/cuda-8.0/lib64/
/usr/local/cuda-8.0
。
此外,請注意,在某些Linux設置中,需要使用nvcc編譯步驟的其他選項。 將-D_MWAITXINTRIN_H_INCLUDED
添加到nvcc命令行以避免來自mwaitxintrin.h
錯誤。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.