c++ - 编译输入为 5d 张量的自定义 tf 操作

Question

我试图通过 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

and here is the register part in the .cc file:

#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));
});

here is the declaration of op class:

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

now comes the cpu operation:

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

    }

at last build registration:

#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

then compile it by g++:

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

**

I got some errors:

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>&}’

**

However, if I change the shape of input and output tensor from 5 to 4, the compile can be successfully done:

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

and

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 维张量。希望有人告诉我如何解决这个问题。谢谢!

Answer 1

看起来您已将 crops 定义为 5D 张量，但您仅使用 4D 索引访问它:crops(b, y, x, d) = extrapolation_value。您可能希望将其作为一个 block 分配给它:https://eigen.tuxfamily.org/dox/group__TutorialBlockOperations.html