25 #include "../TensorFlowHelper.h" 27 #include "tensorflow/core/framework/op.h" 28 #include "tensorflow/core/framework/op_kernel.h" 29 #include "tensorflow/core/lib/core/errors.h" 34 tensorflow::OpKernelConstruction* construction)
35 : OpKernel(construction) {
39 OP_REQUIRES_OK(construction,
40 construction->GetAttr(
"max_hash_table_size",
44 void Compute(tensorflow::OpKernelContext* context)
override {
48 const Tensor&
points = context->input(0);
49 const Tensor& radius = context->input(1);
50 OP_REQUIRES(context, TensorShapeUtils::IsScalar(radius.shape()),
51 errors::InvalidArgument(
"radius must be scalar, got shape ",
52 radius.shape().DebugString()));
54 const Tensor& points_row_splits = context->input(2);
56 const Tensor& hash_table_size_factor_tensor = context->input(3);
59 TensorShapeUtils::IsScalar(
60 hash_table_size_factor_tensor.shape()),
61 errors::InvalidArgument(
62 "hash_table_size_factor must be scalar, got shape ",
63 hash_table_size_factor_tensor.shape().DebugString()));
64 const double hash_table_size_factor =
65 hash_table_size_factor_tensor.scalar<
double>()();
67 Dim num_points(
"num_points");
68 Dim batch_size(
"batch_size");
70 CHECK_SHAPE(context, points_row_splits, batch_size + 1);
72 std::vector<uint32_t> hash_table_splits(batch_size.
value() + 1, 0);
73 for (
int i = 0; i < batch_size.
value(); ++i) {
74 int64_t num_points_i = points_row_splits.flat<int64>()(i + 1) -
75 points_row_splits.flat<int64>()(i);
77 int64_t hash_table_size = std::min<int64_t>(
78 std::max<int64_t>(hash_table_size_factor * num_points_i, 1),
80 hash_table_splits[i + 1] = hash_table_splits[i] + hash_table_size;
83 Tensor* hash_table_index = 0;
84 TensorShape hash_table_index_shape({num_points.
value()});
85 OP_REQUIRES_OK(context,
86 context->allocate_output(0, hash_table_index_shape,
89 Tensor* hash_table_cell_splits = 0;
90 TensorShape hash_table_cell_splits_shape(
91 {hash_table_splits.back() + 1});
92 OP_REQUIRES_OK(context,
93 context->allocate_output(1, hash_table_cell_splits_shape,
94 &hash_table_cell_splits));
96 Tensor* out_hash_table_splits = 0;
97 TensorShape out_hash_table_splits_shape({batch_size.
value() + 1});
98 OP_REQUIRES_OK(context,
99 context->allocate_output(2, out_hash_table_splits_shape,
100 &out_hash_table_splits));
101 for (
size_t i = 0; i < hash_table_splits.size(); ++i) {
102 out_hash_table_splits->flat<
uint32_t>()(i) = hash_table_splits[i];
105 Kernel(context, points, radius, points_row_splits, hash_table_splits,
106 *hash_table_index, *hash_table_cell_splits);
109 virtual void Kernel(tensorflow::OpKernelContext* context,
110 const tensorflow::Tensor&
points,
111 const tensorflow::Tensor& radius,
112 const tensorflow::Tensor& points_row_splits,
113 const std::vector<uint32_t>& hash_table_splits,
114 tensorflow::Tensor& hash_table_index,
115 tensorflow::Tensor& hash_table_cell_splits) = 0;
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c k4a_image_t image_handle uint8_t image_handle image_handle image_handle image_handle uint32_t
Definition: K4aPlugin.cpp:557
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:204
Definition: BuildSpatialHashTableOpKernel.h:31
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:69
Definition: ContinuousConv.h:35
int64_t & value()
Definition: ShapeChecking.h:89
void Compute(tensorflow::OpKernelContext *context) override
Definition: BuildSpatialHashTableOpKernel.h:44
virtual void Kernel(tensorflow::OpKernelContext *context, const tensorflow::Tensor &points, const tensorflow::Tensor &radius, const tensorflow::Tensor &points_row_splits, const std::vector< uint32_t > &hash_table_splits, tensorflow::Tensor &hash_table_index, tensorflow::Tensor &hash_table_cell_splits)=0
int points
Definition: FilePCD.cpp:73
BuildSpatialHashTableOpKernel(tensorflow::OpKernelConstruction *construction)
Definition: BuildSpatialHashTableOpKernel.h:33
int max_hash_table_size
Definition: BuildSpatialHashTableOpKernel.h:118
Definition: ShapeChecking.h:35