26 #include "tensorflow/core/framework/op.h" 27 #include "tensorflow/core/framework/op_kernel.h" 28 #include "tensorflow/core/lib/core/errors.h" 34 template <
class TReal,
class TFeat>
35 class OutputAllocator {
37 OutputAllocator(tensorflow::OpKernelContext* context) : context(context) {}
39 void AllocPooledPositions(TReal** ptr,
size_t num) {
43 TensorShape shape({int64_t(num), 3});
44 OP_REQUIRES_OK(context, context->allocate_output(0, shape, &tensor));
45 auto flat_tensor = tensor->flat<TReal>();
46 *ptr = flat_tensor.data();
49 void AllocPooledFeatures(TFeat** ptr,
size_t num,
int channels) {
53 TensorShape shape({int64_t(num), channels});
54 OP_REQUIRES_OK(context, context->allocate_output(1, shape, &tensor));
55 auto flat_tensor = tensor->flat<TFeat>();
56 *ptr = flat_tensor.data();
60 tensorflow::OpKernelContext* context;
64 class VoxelPoolingGradOpKernel :
public tensorflow::OpKernel {
66 explicit VoxelPoolingGradOpKernel(
67 tensorflow::OpKernelConstruction* construction)
68 : OpKernel(construction) {
71 std::string pos_fn_str;
72 OP_REQUIRES_OK(construction,
73 construction->GetAttr(
"position_fn", &pos_fn_str));
75 if (pos_fn_str ==
"average")
77 else if (pos_fn_str ==
"nearest_neighbor")
82 std::string feat_fn_str;
83 OP_REQUIRES_OK(construction,
84 construction->GetAttr(
"feature_fn", &feat_fn_str));
86 if (feat_fn_str ==
"average")
88 else if (feat_fn_str ==
"nearest_neighbor")
94 void Compute(tensorflow::OpKernelContext* context)
override {
98 const Tensor& positions = context->input(0);
100 context, positions.shape().dims() == 2,
101 errors::InvalidArgument(
"positions must be a rank 2 tensor"));
103 const Tensor& features = context->input(1);
105 context, features.shape().dims() == 2,
106 errors::InvalidArgument(
"features must be a rank 2 tensor"));
108 const Tensor& voxel_size = context->input(2);
110 context, TensorShapeUtils::IsScalar(voxel_size.shape()),
111 errors::InvalidArgument(
"voxel_size must be a scalar, but is ",
112 voxel_size.shape().DebugString()));
114 const Tensor& pooled_positions = context->input(3);
115 OP_REQUIRES(context, pooled_positions.shape().dims() == 2,
116 errors::InvalidArgument(
117 "pooled_positions must be a rank 2 tensor"));
119 const Tensor& pooled_features_gradient = context->input(4);
121 context, pooled_features_gradient.shape().dims() == 2,
122 errors::InvalidArgument(
123 "pooled_features_gradient must be a rank 2 tensor"));
125 Tensor* features_backprop =
nullptr;
126 OP_REQUIRES_OK(context, context->allocate_output(0, features.shape(),
127 &features_backprop));
129 Kernel(context, *features_backprop, positions, features,
130 pooled_positions, pooled_features_gradient, voxel_size);
134 virtual void Kernel(tensorflow::OpKernelContext* context,
135 tensorflow::Tensor& features_backprop,
136 const tensorflow::Tensor& positions,
137 const tensorflow::Tensor& features,
138 const tensorflow::Tensor& pooled_positions,
139 const tensorflow::Tensor& pooled_features_gradient,
140 const tensorflow::Tensor& voxel_size) = 0;
Definition: VoxelPooling.h:40
Definition: VoxelPooling.h:40
Definition: VoxelPooling.h:40
Definition: ContinuousConv.h:35
AccumulationFn
Definition: VoxelPooling.h:40
Definition: VoxelPooling.h:40