编译一个自定义的tf操作,其中输入是5d张量

compile a custom tf operation where input is 5d tensor

本文关键字:输入 张量 5d 操作 tf 一个 自定义 编译      更新时间:2023-10-16

我试图用g++编译一个tensorflow自定义操作,但我遇到了一些错误,我不知道如何解决。这个运算的输入是一个5D张量。这是.h文件

#ifndef TENSORFLOW_CORE_KERNELS_CROP_RESIZE_OP_H_
#define TENSORFLOW_CORE_KERNELS_CROP_RESIZE_OP_H_ 
#include "cuda.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {
namespace functor
{
template <typename Device, typename T>
struct CropResize
{
// We assume that the tensor sizes are correct.
bool operator()(const OpKernelContext* context,
typename TTypes<T, 5>::ConstTensor image,
typename TTypes<float, 2>::ConstTensor boxes,
typename TTypes<int32, 1>::ConstTensor box_ind,
float extrapolation_value,
typename TTypes<float, 5>::Tensor crops);
};
}
}
#endif

,这里是.cc文件中的寄存器部分:

#include <cstring>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "crop_and_resize_op.h"
#include "cuda.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/util/work_sharder.h"
#include "tensorflow/core/util/tensor_format.h"

REGISTER_OP("CropResize")
.Input("image: T")
.Input("boxes: float")
.Input("box_ind: int32")
.Input("crop_size: int32")
.Output("crops: float")
.Attr("T: {uint8, uint16, int8, int16, int32, int64, half, float, double}")
.Attr("method: {'bilinear'} = 'bilinear'")
.Attr("extrapolation_value: float = 0")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
// Get inputs and validate ranks.
ShapeHandle input;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input));
ShapeHandle boxes;
TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &boxes));
ShapeHandle box_ind;
TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &box_ind));
// boxes[0] and box_ind[0] are both num_boxes.
DimensionHandle num_boxes_dim;
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(boxes, 0), c->Dim(box_ind, 0), &num_boxes_dim));
// boxes.dim(1) is 4.
DimensionHandle unused;
TF_RETURN_IF_ERROR(c->WithValue(c->Dim(boxes, 1), 6, &unused));
return SetOutputToSizedImage(c, num_boxes_dim, 3 /* size_input_idx */,
c->Dim(input, 4));
});

这里是op类的声明:

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;
namespace {
static inline Status ParseAndCheckBoxSizes(const Tensor& boxes, const Tensor& box_index, int* num_boxes)
{
if (boxes.NumElements() == 0 && box_index.NumElements() == 0) {
*num_boxes = 0;
return Status::OK();
}
// The shape of 'boxes' is [num_boxes, 6].
if (boxes.dims() != 2) {
return errors::InvalidArgument("boxes must be 2-D",
boxes.shape().DebugString());
}
*num_boxes = boxes.dim_size(0);
if (boxes.dim_size(1) != 6) {
return errors::InvalidArgument("boxes must have 6 columns");
}
// The shape of 'box_index' is [num_boxes].
if (box_index.dims() != 1) {
return errors::InvalidArgument("box_index must be 1-D",
box_index.shape().DebugString());
}
if (box_index.dim_size(0) != *num_boxes) {
return errors::InvalidArgument("box_index has incompatible shape");
}
return Status::OK();
}
}
template <typename Device, typename T>
class CropResizeOp : public OpKernel {
public:
explicit CropResizeOp(OpKernelConstruction* context)
: OpKernel(context) {
string method;
OP_REQUIRES_OK(context, context->GetAttr("method", &method));
OP_REQUIRES(context, method == "bilinear",
errors::InvalidArgument("method must be 'bilinear'", method));
OP_REQUIRES_OK(context, context->GetAttr("extrapolation_value",
&extrapolation_value_));
}
void Compute(OpKernelContext* context) override {
// The shape of 'image' is [batch_size, image_height, image_width, image_depth,
// channels].
const Tensor& image = context->input(0);
// The shape of 'boxes' is [num_boxes, 6].
const Tensor& boxes = context->input(1);
// The shape of 'box_index' is [num_boxes].
const Tensor& box_index = context->input(2);
// The shape of 'crop_size' is [3].
const Tensor& crop_size = context->input(3);
// Validate inputs dimensions.
OP_REQUIRES(context, image.dims() == 5,
errors::InvalidArgument("input image must be 5-D",
image.shape().DebugString()));
const int batch_size = image.dim_size(0);
const int image_height = image.dim_size(1);
const int image_width = image.dim_size(2);
const int image_depth = image.dim_size(3);
const int depth = image.dim_size(4);
OP_REQUIRES(
context, image_height > 0 && image_width > 0,
errors::InvalidArgument("image dimensions must be positive"));
int num_boxes = 0;
OP_REQUIRES_OK(
context, ParseAndCheckBoxSizes(boxes, box_index, &num_boxes));
OP_REQUIRES(context, crop_size.dims() == 1,
errors::InvalidArgument("crop_size must be 1-D",
crop_size.shape().DebugString()));
OP_REQUIRES(
context, crop_size.dim_size(0) == 3,
errors::InvalidArgument("crop_size must have three elements",
crop_size.shape().DebugString()));
// Copy and validate crop sizes.
auto crop_size_vec = crop_size.vec<int32>();
//          const int crop_height = ::tensorflow::internal::SubtleMustCopy(crop_size_vec(0));
//          const int crop_width = ::tensorflow::internal::SubtleMustCopy(crop_size_vec(1));
const int crop_height = crop_size_vec(0);
const int crop_width = crop_size_vec(1);
const int crop_depth = crop_size_vec(2);
OP_REQUIRES(
context, crop_height > 0 && crop_width > 0 && crop_depth > 0,
errors::InvalidArgument("crop dimensions must be positive"));
// Allocate output tensor.
Tensor* output = nullptr;
OP_REQUIRES_OK(
context,
context->allocate_output(
0, TensorShape({ num_boxes, crop_height, crop_width, crop_depth, depth }),
&output));

const bool status = functor::CropResize<Device, T>()(
context, image.tensor<T,5>(), boxes.tensor<float, 2>(),
box_index.tensor<int32, 1>(), extrapolation_value_,
output->tensor<float,5>());
if (!status) {
context->SetStatus(
errors::Internal("Failed launch CropAndResizeKernel."));
}
}
private:
float extrapolation_value_;
};

现在是cpu操作:

namespace functor {
template <typename T>
struct CropResize<CPUDevice, T> {
bool operator()(const OpKernelContext* context,
typename TTypes<T,5>::ConstTensor image,
typename TTypes<float, 2>::ConstTensor boxes,
typename TTypes<int32, 1>::ConstTensor box_index,
float extrapolation_value,
typename TTypes<float,5>::Tensor crops) {
const int batch_size = image.dimension(0);
const int image_height = image.dimension(1);
const int image_width = image.dimension(2);
const int num_boxes = crops.dimension(0);
const int crop_height = crops.dimension(1);
const int crop_width = crops.dimension(2);
const int depth = crops.dimension(3);
// Sharding across boxes.
//auto CropAndResizePerBox = [&](int start_box, int limit_box) {
//for (int b = start_box; b < limit_box; ++b) {
for (int b = 0; b < num_boxes; ++b) {
const float y1 = boxes(b, 0);
const float x1 = boxes(b, 1);
const float y2 = boxes(b, 2);
const float x2 = boxes(b, 3);
const int32 b_in = box_index(b);
//                      if (!FastBoundsCheck(b_in, batch_size)) {
//                          continue;
//                      }
const float height_scale =
(crop_height > 1)
? (y2 - y1) * (image_height - 1) / (crop_height - 1)
: 0;
const float width_scale =
(crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1)
: 0;
for (int y = 0; y < crop_height; ++y) {
const float in_y = (crop_height > 1)
? y1 * (image_height - 1) + y * height_scale
: 0.5 * (y1 + y2) * (image_height - 1);
if (in_y < 0 || in_y > image_height - 1) {
for (int x = 0; x < crop_width; ++x) {
for (int d = 0; d < depth; ++d) {
crops(b, y, x, d) = extrapolation_value;
}
}
continue;
}
const int top_y_index = floorf(in_y);
const int bottom_y_index = ceilf(in_y);
const float y_lerp = in_y - top_y_index;
for (int x = 0; x < crop_width; ++x) {
const float in_x = (crop_width > 1)
? x1 * (image_width - 1) + x * width_scale
: 0.5 * (x1 + x2) * (image_width - 1);
if (in_x < 0 || in_x > image_width - 1) {
for (int d = 0; d < depth; ++d) {
crops(b, y, x, d) = extrapolation_value;
}
continue;
}
const int left_x_index = floorf(in_x);
const int right_x_index = ceilf(in_x);
const float x_lerp = in_x - left_x_index;
for (int d = 0; d < depth; ++d) {
const float top_left(static_cast<float>(
image(b_in, top_y_index, left_x_index, d)));
const float top_right(static_cast<float>(
image(b_in, top_y_index, right_x_index, d)));
const float bottom_left(static_cast<float>(
image(b_in, bottom_y_index, left_x_index, d)));
const float bottom_right(static_cast<float>(
image(b_in, bottom_y_index, right_x_index, d)));
const float top = top_left + (top_right - top_left) * x_lerp;
const float bottom =
bottom_left + (bottom_right - bottom_left) * x_lerp;
crops(b, y, x, d) = top + (bottom - top) * y_lerp;
}
}
}
};
return true;
}
};
}

在上次生成注册时:

#define REGISTER_KERNEL(T)                                
REGISTER_KERNEL_BUILDER(Name("CropResize")           
.Device(DEVICE_CPU)         
.TypeConstraint<T>("T")     
.HostMemory("crop_size"),   
CropResizeOp<CPUDevice, T>);      
TF_CALL_float(REGISTER_KERNEL);
//TF_CALL_double(REGISTER_KERNEL);

//TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);
#undef REGISTER_KERNEL

然后通过g++编译:

g++ -std=c++11 -shared crop_and_resize_op.cc -o crop_and_resize_op.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

**

我得到了一些错误:

crop_and_resize_op.cc:244:56:   required from ‘void tensorflow::CropResizeOp<Device, T>::Compute(tensorflow::OpKernelContext*) [with Device = Eigen::ThreadPoolDevice; T = float]’
crop_and_resize_op.cc:772:1:   required from here
/usr/local/lib/python3.5/dist-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorMap.h:239:7: **error**: static assertion failed: Number of indices used to access a tensor coefficient must be equal to the rank of the tensor.
static_assert(sizeof...(otherIndices) + 2 == NumIndices || NumIndices == Dynamic, "Number of indices used to access a tensor coefficient must be equal to the rank of the tensor.");
^
/usr/local/lib/python3.5/dist-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorMap.h:242:123: **error**: no matching function for call to ‘Eigen::DSizes<long int, 5>::IndexOfRowMajor(Eigen::array<long int, 4ul>)’
const Index index = m_dimensions.IndexOfRowMajor(array<Index, NumDims>{{firstIndex, secondIndex, otherIndices...}});
                                                     ^
In file included from /usr/local/lib/python3.5/dist-packages/tensorflow/include/unsupported/Eigen/CXX11/Tensor:102:0,
from /usr/local/lib/python3.5/dist-packages/tensorflow/include/third_party/eigen3/unsupported/Eigen/CXX11/Tensor:1,
from crop_and_resize_op.cc:2:
/usr/local/lib/python3.5/dist-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorDimensions.h:330:52: note: candidate: DenseIndex Eigen::DSizes<DenseIndex, NumDims>::IndexOfRowMajor(Eigen::array<DenseIndex, NumDims>&) const [with DenseIndex = long int; int NumDims = 5; Eigen::array<DenseIndex, NumDims> = std::array<long int, 5ul>]
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DenseIndex IndexOfRowMajor(const array<DenseIndex, NumDims>& indices) const {
^
/usr/local/lib/python3.5/dist-packages/tensorflow/include/unsupported/Eigen/CXX11/src/Tensor/TensorDimensions.h:330:52: note:   no known conversion for argument 1 from ‘Eigen::array<long int, 4ul> {aka std::array<long int, 4ul>}’ to ‘Eigen::array<long int, 5ul>& {aka const std::array<long int, 5ul>&}’

**

然而,如果我将输入和输出张量的形状从5改为4,编译可以成功完成:

const bool status = functor::CropResize<Device, T>()(
context, image.tensor<T,4>(), boxes.tensor<float, 2>(),
box_index.tensor<int32, 1>(), extrapolation_value_,
output->tensor<float,4>());

bool operator()(const OpKernelContext* context,
typename TTypes<T,4>::ConstTensor image,
typename TTypes<float, 2>::ConstTensor boxes,
typename TTypes<int32, 1>::ConstTensor box_index,
float extrapolation_value,
typename TTypes<float,4>::Tensor crops) 

我不知道这是怎么来的,但我确实需要让输入和输出是一个5-D张量。希望有人告诉我如何解决这个问题。非常感谢。

看起来您已经将crops定义为5D张量,但您只使用4D索引访问它:crops(b, y, x, d) = extrapolation_value。您可能希望将其指定为块:https://eigen.tuxfamily.org/dox/group__TutorialBlockOperations.html