Open3D (C++ API)  0.18.0+252c867
SparseConvTranspose.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // Copyright (c) 2018-2023 www.open3d.org
5 // SPDX-License-Identifier: MIT
6 // ----------------------------------------------------------------------------
7 
8 #pragma once
9 #include <tbb/parallel_for.h>
10 
11 #include <Eigen/Core>
12 
13 namespace open3d {
14 namespace ml {
15 namespace impl {
16 
19 template <class TFeat,
20  class TOut,
21  class TIndex,
22  class TKernelIndex,
23  bool NORMALIZE>
25  TOut* out_features,
26  const std::vector<int>& filter_dims,
27  const TFeat* filter,
28  size_t num_out,
29  const TFeat* out_importance,
30  size_t num_inp,
31  const TFeat* inp_features,
32  const TFeat* inp_neighbors_importance_sum,
33  const int64_t* inp_neighbors_row_splits,
34  const TIndex* neighbor_index,
35  const TKernelIndex* neighbors_kernel_index,
36  const TFeat* neighbor_importance,
37  const int64_t* neighbors_row_splits) {
38  const bool NEIGHBOR_IMPORTANCE = inp_neighbors_importance_sum;
39 
40  const int in_channels = filter_dims[filter_dims.size() - 2];
41  const int out_channels = filter_dims[filter_dims.size() - 1];
42 
43  int num_kernel_elements = 1;
44  for (int i = 0; i < filter_dims.size() - 2; ++i)
45  num_kernel_elements *= filter_dims[i];
46 
47  memset(out_features, 0, sizeof(TOut) * num_out * out_channels);
48 
49  tbb::parallel_for(
50  tbb::blocked_range<size_t>(0, num_out, 32),
51  [&](const tbb::blocked_range<size_t>& r) {
52  int range_length = r.end() - r.begin();
53 
54  Eigen::Map<Eigen::Matrix<TOut, Eigen::Dynamic, Eigen::Dynamic>>
55  C(out_features + (r.begin() * out_channels),
56  out_channels, range_length);
57 
58  for (size_t out_idx = r.begin(); out_idx != r.end();
59  ++out_idx) {
60  const int out_col = out_idx - r.begin();
61  const size_t neighbor_start = neighbors_row_splits[out_idx];
62  const size_t neighbor_end =
63  neighbors_row_splits[out_idx + 1];
64 
65  for (size_t n = neighbor_start; n < neighbor_end; ++n) {
66  const size_t inp_idx = neighbor_index[n];
67  const int kernel_idx = neighbors_kernel_index[n];
68 
69  TFeat n_importance = NEIGHBOR_IMPORTANCE
70  ? neighbor_importance[n]
71  : TFeat(1);
72 
73  TFeat normalizer(1);
74  if (NORMALIZE) {
75  if (NEIGHBOR_IMPORTANCE) {
76  if (inp_neighbors_importance_sum[inp_idx] !=
77  TFeat(0))
78  normalizer /= inp_neighbors_importance_sum
79  [inp_idx];
80  } else {
81  size_t num_inp_neighbors;
82  const size_t inp_neighbor_start =
83  inp_neighbors_row_splits[inp_idx];
84  const size_t inp_neighbor_end =
85  inp_neighbors_row_splits[inp_idx + 1];
86  num_inp_neighbors =
87  inp_neighbor_end - inp_neighbor_start;
88  if (num_inp_neighbors > 0)
89  normalizer /= TFeat(num_inp_neighbors);
90  }
91  }
92 
93  Eigen::Map<const Eigen::Matrix<TFeat, Eigen::Dynamic,
94  Eigen::Dynamic>>
95  A(filter + kernel_idx * out_channels *
96  in_channels,
97  out_channels, in_channels);
98 
99  Eigen::Map<const Eigen::Matrix<TFeat, Eigen::Dynamic,
100  Eigen::Dynamic>>
101  B(inp_features + inp_idx * in_channels,
102  in_channels, 1);
103  TFeat scale = normalizer * n_importance;
104  C.col(out_col) +=
105  (A * (scale * B)).template cast<TOut>();
106  }
107 
108  } // out_idx
109 
110  if (out_importance) {
111  for (int i = 0; i < range_length; ++i)
112  C.col(i) *= TOut(out_importance[r.begin() + i]);
113  }
114  });
115 }
116 
163 template <class TFeat, class TOut, class TIndex, class TKernelIndex>
165  TOut* out_features,
166  const std::vector<int>& filter_dims,
167  const TFeat* filter,
168  size_t num_out,
169  const TFeat* out_importance,
170  size_t num_inp,
171  const TFeat* inp_features,
172  const TFeat* inp_neighbors_importance_sum,
173  const int64_t* inp_neighbors_row_splits,
174  const TIndex* neighbor_index,
175  const TKernelIndex* neighbors_kernel_index,
176  const TFeat* neighbor_importance,
177  const int64_t* neighbors_row_splits,
178  bool normalize) {
179 #define FN_PARAMETERS \
180  out_features, filter_dims, filter, num_out, out_importance, num_inp, \
181  inp_features, inp_neighbors_importance_sum, \
182  inp_neighbors_row_splits, neighbor_index, neighbors_kernel_index, \
183  neighbor_importance, neighbors_row_splits
184 
185 #define CALL_TEMPLATE(NORMALIZE) \
186  if (NORMALIZE == normalize) \
187  _SparseConvTransposeComputeFeaturesCPU<TFeat, TOut, TIndex, \
188  TKernelIndex, NORMALIZE>( \
189  FN_PARAMETERS);
190 
191 #define CALL_TEMPLATE2 \
192  CALL_TEMPLATE(true) \
193  CALL_TEMPLATE(false)
194 
196 
197 #undef CALL_TEMPLATE
198 #undef CALL_TEMPLATE2
199 
200 #undef FN_PARAMETERS
201 }
202 
203 } // namespace impl
204 } // namespace ml
205 } // namespace open3d
#define CALL_TEMPLATE2
void _SparseConvTransposeComputeFeaturesCPU(TOut *out_features, const std::vector< int > &filter_dims, const TFeat *filter, size_t num_out, const TFeat *out_importance, size_t num_inp, const TFeat *inp_features, const TFeat *inp_neighbors_importance_sum, const int64_t *inp_neighbors_row_splits, const TIndex *neighbor_index, const TKernelIndex *neighbors_kernel_index, const TFeat *neighbor_importance, const int64_t *neighbors_row_splits)
Definition: SparseConvTranspose.h:24
void SparseConvTransposeComputeFeaturesCPU(TOut *out_features, const std::vector< int > &filter_dims, const TFeat *filter, size_t num_out, const TFeat *out_importance, size_t num_inp, const TFeat *inp_features, const TFeat *inp_neighbors_importance_sum, const int64_t *inp_neighbors_row_splits, const TIndex *neighbor_index, const TKernelIndex *neighbors_kernel_index, const TFeat *neighbor_importance, const int64_t *neighbors_row_splits, bool normalize)
Definition: SparseConvTranspose.h:164
Definition: PinholeCameraIntrinsic.cpp:16