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 VoxelPoolingOpKernel :
public tensorflow::OpKernel {
66 explicit VoxelPoolingOpKernel(
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")
93 OP_REQUIRES_OK(construction, construction->GetAttr(
"debug", &debug));
96 void Compute(tensorflow::OpKernelContext* context)
override {
99 const Tensor& positions = context->input(0);
101 context, positions.shape().dims() == 2,
102 errors::InvalidArgument(
"positions must be a rank 2 tensor"));
104 const Tensor& features = context->input(1);
106 context, features.shape().dims() == 2,
107 errors::InvalidArgument(
"features must be a rank 2 tensor"));
109 const Tensor& voxel_size = context->input(2);
111 context, TensorShapeUtils::IsScalar(voxel_size.shape()),
112 errors::InvalidArgument(
"voxel_size must be a scalar, but is ",
113 voxel_size.shape().DebugString()));
115 Kernel(context, positions, features, voxel_size);
119 virtual void Kernel(tensorflow::OpKernelContext* context,
120 const tensorflow::Tensor& positions,
121 const tensorflow::Tensor& features,
122 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