Ajout d'une opération GPU dans Tensorflow

Question

J'essaye d'ajouter une nouvelle op à TensorFlow suivant docilement this. La différence étant que j'essaie de mettre en œuvre une opération basée sur GPU. L'op que j'essaye d'ajouter est l'op cuda de here (cuda_op.py, cuda_op_kernel.cc, cuda_op_kernel.cu.cc). . Je suis en train de compiler ces dehors de tensorflow et l'utilisation tf.load_op_library pour les tirer dans je l'ai fait quelques changements alors voici mes fichiers:

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

Je suis en mesure de compiler et exécuter avec succès test.py, mais la sortie est toujours [0 0 0 0 0]. Si je remplace AddOneKernel<<<32, 256>>>(in, N, out); avec for (int i = 0; i < N; i++) out[i] = in[i] + 1; et DEVICE_GPU avec DEVICE_CPU, l'op produit les bonnes valeurs [6 5 4 3 2] (avec exactement le même CMakeList.txt).

Une idée de comment obtenir les bonnes valeurs à retourner?

Answer 1

Je ne me rappelle pas exactement où j'ai trouvé les trucs de cmake pour CUDA, mais les options perturbaient la compilation en quelque sorte. Remplacement cuda_add_library dans le CMakeLists.txt pour être le suivant a résolu le problème.

#no options needed 
cuda_add_library(
    cuda_op SHARED 
    src/cuda_op_kernel.cu 
    src/cuda_op_kernel.cc) 

Answer 2

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

- Configuration fait

- Génération fait

- Construire des fichiers ont été par écrit à:/home/ubuntu/Desktop/src/src/build

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

[33%] Bâtiment FPUNV objet (Device) CMakeFiles/cuda_op.d/cuda_op_generated_cuda_op_kernel.cu.o

avertissement

de nvcc: The 'compute_20', 'sm_20', et les architectures 'sm_21' sont obsolètes, et peuvent être supprimé dans une version ultérieure (Utilisez -Wno-deprecated-gpu-targets pour supprimer l'avertissement). Nvcc warning: Les architectures 'compute_20', 'sm_20' et 'sm_21' sont obsolètes et peuvent être supprimées dans une version ultérieure (utilisez -Wno-deprecated-gpu-targets pour supprimer l'avertissement).

dépendances d'analyse de cible cuda_op

[66%] objet de construction CXX CMakeFiles/cuda_op.dir/cuda_op_kernel.cc.o /home/ubuntu/Desktop/src/src/cuda_op_kernel.cc: 1: 17: erreur: 'tensorflow' n'est pas un nom d'espace de noms utilisant l'espace de noms tensorflow; // NOLINT (build/namespaces)