Question

我正在使用Tensorflow 1.9rc0，它是从使用bazel版本0.15.0的源编译而来的。我正在将cuda支持与cuda 9.2版和cudnn 7版一起使用。

我正在尝试构建一个执行cuda内核的自定义操作。我已遵循doumentation的相关规定，并检查了一些已实施的操作来进行开发。

我最终得到了以下代码：

kernel_example.h：

#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/tensor_types.h"

#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/platform/types.h"


namespace functor {

template <typename Device,typename T>
struct ExampleFunctor {
  void operator()(const Device& d,
                  const T* input, const T* filter, T* output,
                  int in_depth, int input_cols, int input_rows,
                  int out_depth, int filter_cols, int filter_rows,
                  int stride_rows, int stride_cols,
                  int n_elements);
};

#if GOOGLE_CUDA

typedef Eigen::GpuDevice GPUDevice;

template <typename T>
struct ExampleFunctor<GPUDevice, T> {
    void operator()(
        const GPUDevice& d,
        const T* input, const T* filter, T* output,
        int in_depth, int input_cols, int input_rows,
        int out_depth, int filter_cols, int filter_rows,
        int stride_rows, int stride_cols,
        int n_elements);
};

#endif //GOOGLE_CUDA

}

kernel_example.cc：

#include "kernel_example.h"
#include "tensorflow/core/framework/op_kernel.h"

#include "tensorflow/core/framework/common_shape_fns.h"

using namespace tensorflow;

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

REGISTER_OP("MyConvGpu")
    .Input("input: T")
    .Input("filter: T")
    .Output("output: T")
    .SetShapeFn(tensorflow::shape_inference::UnknownShape);

// OpKernel definition.
// template parameter <T> is the datatype of the tensors.
template <typename Device, typename T>
class ExampleOp : public OpKernel {
 public:
  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Loading op parameters and defining  variables

    functor::ExampleFunctor<Device, T> functor;
    functor(
        context->eigen_device<Device>(),
        input.flat<T>().data(),
        filter.flat<T>().data(),
        output->flat<T>().data(),
        in_depth, input_cols, input_rows,
        out_depth, filter_cols, filter_rows,
        stride_rows, stride_cols,
        static_cast<int>(output->NumElements()));
  }
};

#if GOOGLE_CUDA

#define REGISTER_GPU_KERNEL(T)                                    \
  REGISTER_KERNEL_BUILDER(Name("Example")                         \
                              .Device(DEVICE_GPU)                 \
                              .TypeConstraint<T>("T"),            \
                              ExampleOp<GPUDevice, T>);

TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL);

//REGISTER_GPU_KERNEL(Eigen::half)
REGISTER_GPU_KERNEL(float)
REGISTER_GPU_KERNEL(double)
REGISTER_GPU_KERNEL(int)

#endif

kernel_example.cu.cc：

#ifdef GOOGLE_CUDA

#define EIGEN_USE_GPU
#include "kernel_example.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"

using namespace tensorflow;

namespace functor {

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const T* input, const T* filter, T* output,
                               const int batch_size,
                               const int in_depth, const int input_cols, const int input_rows,
                               const int out_depth, const int filter_cols, const int filter_rows,
                               const int stride_rows, const int stride_cols,
                               const int n_elements) {
    // Kernel here
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const Eigen::GpuDevice& d,
    const T* input, const T* filter, T* output,
    int in_depth, int input_cols, int input_rows,
    int out_depth, int filter_cols, int filter_rows,
    int stride_rows, int stride_cols,
    int n_elements) {

  // Launch the cuda kernel.
  ExampleCudaKernel<T>
        <<<(n_elements + 255) / 256, 256>>>(input, filter, output,
                                             batch_size,
                                             in_depth, input_cols, input_rows,
                                             out_depth, filter_cols, filter_rows,
                                             stride_rows, stride_cols,
                                             n_elements);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, double>;
template struct ExampleFunctor<GPUDevice, int>;

} //namespace functor

#endif  // GOOGLE_CUDA

一切正常编译，生成一个 example.so 库文件，并带有以下bazel构建文件：

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "example.so",
    srcs = ["kernel_example.h", "kernel_example.cc"],
    gpu_srcs = ["kernel_example.h", "kernel_example.cu.cc"],
)

但是，使用

将此库加载到Tensorflow执行中时

module = tf.load_op_library('./example.so')

我得到以下输出：

Traceback (most recent call last):
  File "mnist.py", line 51, in <module>
    my_conv_gpu_module = tf.load_op_library('./example.so')
  File "/usr/lib/python3.6/site-packages/tensorflow/python/framework/load_library.py", line 56, in load_op_library
    lib_handle = py_tf.TF_LoadLibrary(library_filename)
tensorflow.python.framework.errors_impl.NotFoundError: ./example.so: undefined symbol: _ZN10tensorflow16FormatFromStringERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEEPNS_12TensorFormatE

我已经开发了其他不使用cuda加速的操作，即使它们的实现非常相似，加载它们也没有问题。

此外，我已经阅读了有关此错误的其他主题，但是该解决方案似乎总是将-cxxopt =“-D_GLIBCXX_USE_CXX11_ABI = 0” 标志添加到bazel构建参数中，因为Tensoflow的二进制文件pip软件包是使用gcc版本4构建的。我正在这样做，但错误仍然存在。

我还在Tensorflow rc1.8，cuda 8和cudnn 6的不同环境下尝试了此代码。

我想念什么？

加载自定义GPU op时的Tensorflow“未定义符号”

0 个答案: