问题描述
我跟随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
对象的形状,以便可以将其作为函数的另一个参数?
解决方法
暂无找到可以解决该程序问题的有效方法,小编努力寻找整理中!
如果你已经找到好的解决方法,欢迎将解决方案带上本链接一起发送给小编。
小编邮箱:dio#foxmail.com (将#修改为@)