tensorflow自定义op实现，如何从Eigen :: TensorMap对象获取类似数据的3D数组？

如何解决tensorflow自定义op实现，如何从Eigen :: TensorMap对象获取类似数据的3D数组？

我跟随https://www.programmersought.com/article/52372107537/最终成功地编译了一个自定义操作，以便在Ubuntu 16.04上以tf = 1.15，cuda = 8.0，cudnn = 7在GPU和CPU上执行AddOne。但是，现在我想更改op，使其做一些更复杂的操作，例如，简单的图像模糊op（https://docs.opencv.org/master/d4/d86/group__imgproc__filter.html#ga8c45db9afe636703801b0b2e440fce37）。 这是我原始的multi_op_kernel.cc代码，用于实现和注册AddOne op：

#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#define EIGEN_USE_THREADS
#endif  // GOOGLE_CUDA
#include "multi_op_kernel.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;

REGISTER_OP("AddOne")
    .Input("input: int32")
    .Output("output: int32")
    .Doc(R"doc(
Adds 1 to all elements of the tensor.
output: A Tensor.
  output = input + 1
)doc");

// Explicit specialization
template<> void AddOneKernelLauncher(const CPUDevice& d,int size,const int* in,int* out){
    for (int i = 0; i < size; i++){
        out[i] = in[i] + 1;
    }
}


template<typename Device>
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();

    AddOneKernelLauncher<Device>(
        context->eigen_device<Device>(),static_cast<int>(input_tensor.NumElements()),input_tensor.flat<int>().data(),output_tensor->flat<int>().data());

  }
};

// Register Kernel for CPU and GPU
REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_CPU),AddOneOp<CPUDevice>);
REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU),AddOneOp<GPUDevice>);

还有头文件multi_op_kernel.h：

// kernel_example.h
#ifndef KERNEL_MULTI_OP_KERNEL_H_
#define KERNEL_MULTI_OP_KERNEL_H_

template <typename Device>
void AddOneKernelLauncher(const Device& d,int* out);

#endif //KERNEL_MULTI_OP_KERNEL_H_

和multi_op_kernel.cu.cc：

#if GOOGLE_CUDA
#define EIGEN_USE_GPU

#include "multi_op_kernel.h"
#include "eigen3/unsupported/Eigen/CXX11/Tensor"

typedef Eigen::GpuDevice GPUDevice;

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

// Explicit specialization
template<> void AddOneKernelLauncher(const GPUDevice& d,int* out) {
  AddOneKernel<<<32,256>>>(in,size,out);
}

// Explicit template instantiation
template void AddOneKernelLauncher<GPUDevice>(const GPUDevice&,int,const int*,int*);
#endif

这是python测试脚本：

import tensorflow as tf


file = './multi_op_kernel.so'
cuda_op_module = tf.load_op_library(file)
config = tf.ConfigProto(log_device_placement=True)
with tf.device("/cpu:0"):
    inp_cpu = tf.constant([[6,4],[2,4]])
    y_cpu = cuda_op_module.add_one(inp_cpu)
with tf.device("/gpu:0"):
    inp_gpu = tf.constant([[6,4]])
    y_gpu = cuda_op_module.add_one(inp_gpu)
with tf.Session(config=config) as sess:
    input_cpu,input_gpu,output_cpu,output_gpu = sess.run([inp_cpu,inp_gpu,y_cpu,y_gpu])

print(' = ' * 10,'CPU',' = ' * 10)
print('Input is: \n',input_cpu)
print('Output is: \n',output_cpu)

print(' = ' * 10,'GPU',input_gpu)
print('Output is: \n',output_gpu)

请首先查看.cc文件。在compute方法中，我将调用AddOneKernelLauncher函数进行实际计算，其中一个实现用于CPU，另一个实现用于.cu.cc文件上的GPU。 我传递给该函数的是4个参数：Device，int size，int* in和int* out。 我发现input_tensor.flat<int>()返回一个Eigen::TensorMap对象，而.data()方法将返回一个const int*类型的数据。但是我只能使用in[i]运算符对数据进行切片。我该怎么办image_blur op。例如在数据的三个维度上？ 例如，（在[i-1] [j-1] [k-1] + in [i-1] [j-1] [k] + in [i-1] [j-1] [ k + 1] + in [i] [j] [k-1] + in [i] [j] [k] + in [i] [j] [k + 1] + in [i + 11] [j +1] [k-1] + in [i + 1] [j + 1] [k] + in [i + 1] [j + 1] [k + 1]）/ 9 ，像这样Eigen::TensorMap对象进行操作？或Eigen::TensorMap.data()对象？ 看来Eigen::TensorMap.data()仅返回类似展平的一维数组的东西。我想要类似3D的东西。

实际上我对C ++还是很陌生，我以前经常用python编程。

编辑： 我到目前为止尝试过的方法如下，看起来很愚蠢...

// Explicit specialization
template<> void AddOneKernelLauncher(const CPUDevice& d,int* out){
    int in_3d[2][2];
    int cnt = 0;
    for (int i = 0; i < 2; i++){
        for (int j = 0; j < 2; j++){
            in_3d[i][j] = in[cnt++];
        }
    }
    int out_3d[2][2];
    for (int i = 0; i < 2; i++){
        for (int j = 0; j < 2; j++){
            out_3d[i][j] = in_3d[j][i];
        }
    }
    int cnt2 = 0;
    for (int i = 0; i < 2; i++){
        for (int j = 0; j < 2; j++){
            out[cnt2++] = out_3d[i][j];
        }
    }
}

这里，我只是做一个实验，将大小为4的1-d数组转换为大小为2x2的2-d数组，并将其转置，然后将其展平为最终输出。但这看起来很幼稚...另外如何获得Eigen::TensorMap对象的形状，以便可以将其作为函数的另一个参数？