SparseConvBackpropFilter.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
#include <tbb/parallel_for.h>
 
#include <Eigen/Core>
#include <mutex>
 
namespace open3d {
namespace ml {
namespace impl {
 
// Implementation of SparseConvBackropFilterCPU
template <class TFeat,
          class TOut,
          class TIndex,
          class TKernelIndex,
          bool POINT_IMPORTANCE>
void _SparseConvBackropFilterCPU(TOut* filter_backprop,
                                 const std::vector<int>& filter_dims,
                                 size_t num_out,
                                 size_t num_inp,
                                 const TFeat* inp_features,
                                 const TFeat* inp_importance,
                                 const TIndex* neighbors_index,
                                 const TKernelIndex* neighbors_kernel_index,
                                 const TFeat* neighbors_importance,
                                 const int64_t* neighbors_row_splits,
                                 const TFeat* out_features_gradient,
                                 bool normalize) {
    const bool NEIGHBOR_IMPORTANCE = neighbors_importance;
 
    const int in_channels = filter_dims[filter_dims.size() - 2];
    const int out_channels = filter_dims[filter_dims.size() - 1];
 
    int num_kernel_elements = 1;
    for (int i = 0; i < filter_dims.size() - 2; ++i)
        num_kernel_elements *= filter_dims[i];
    const int total_filter_size =
            num_kernel_elements * in_channels * out_channels;
 
    memset(filter_backprop, 0, sizeof(TOut) * total_filter_size);
    std::mutex filter_backprop_mutex;
 
    tbb::parallel_for(
            tbb::blocked_range<size_t>(0, num_out, 10032),
            [&](const tbb::blocked_range<size_t>& r) {
                int range_length = r.end() - r.begin();
 
                Eigen::Matrix<TFeat, Eigen::Dynamic, Eigen::Dynamic> B(
                        in_channels * num_kernel_elements, range_length);
                B.setZero();
                Eigen::Matrix<TFeat, Eigen::Dynamic, Eigen::Dynamic> C(
                        out_channels, range_length);
 
                Eigen::Array<TFeat, Eigen::Dynamic, 1> infeat(in_channels, 1);
 
                for (size_t out_idx = r.begin(); out_idx != r.end();
                     ++out_idx) {
                    const int out_col = out_idx - r.begin();
                    const size_t neighbor_start = neighbors_row_splits[out_idx];
                    const size_t neighbor_end =
                            neighbors_row_splits[out_idx + 1];
                    TFeat normalizer(0);
 
                    for (size_t n = neighbor_start; n < neighbor_end; ++n) {
                        const size_t inp_idx = neighbors_index[n];
                        const int kernel_idx = neighbors_kernel_index[n];
 
                        const TFeat n_importance =
                                (NEIGHBOR_IMPORTANCE ? neighbors_importance[n]
                                                     : TFeat(1));
                        normalizer += n_importance;
 
                        for (int ic = 0; ic < in_channels; ++ic)
                            infeat(ic) =
                                    inp_features[inp_idx * in_channels + ic];
 
                        TFeat importance(1);
                        if (POINT_IMPORTANCE)
                            importance = inp_importance[inp_idx];
                        if (NEIGHBOR_IMPORTANCE) importance *= n_importance;
 
                        if (POINT_IMPORTANCE || NEIGHBOR_IMPORTANCE) {
                            for (int ic = 0; ic < in_channels; ++ic)
                                infeat(ic) *= importance;
                        }
                        for (int ic = 0; ic < in_channels; ++ic) {
                            B(kernel_idx * in_channels + ic, out_col) =
                                    infeat(ic);
                        }
                    }
 
                    C.col(out_col) = Eigen::Map<
                            const Eigen::Array<TFeat, Eigen::Dynamic, 1>>(
                            out_features_gradient + out_idx * out_channels,
                            out_channels, 1);
 
                    if (normalize && normalizer != TFeat(0))
                        C.col(out_col) /= normalizer;
 
                }  // out_idx
 
                Eigen::Matrix<TFeat, Eigen::Dynamic, Eigen::Dynamic> A(
                        out_channels, num_kernel_elements * in_channels);
 
                A = C * B.transpose();
 
                {
                    std::lock_guard<std::mutex> lock(filter_backprop_mutex);
                    int linear_i = 0;
                    for (int j = 0; j < num_kernel_elements * in_channels; ++j)
                        for (int i = 0; i < out_channels; ++i, ++linear_i) {
                            filter_backprop[linear_i] += TOut(A(i, j));
                        }
                }
            });
}
 
template <class TFeat, class TOut, class TIndex, class TKernelIndex>
void SparseConvBackpropFilterCPU(TOut* filter_backprop,
                                 const std::vector<int>& filter_dims,
                                 size_t num_out,
                                 size_t num_inp,
                                 const TFeat* inp_features,
                                 const TFeat* inp_importance,
                                 const TIndex* neighbors_index,
                                 const TKernelIndex* neighbors_kernel_index,
                                 const TFeat* neighbors_importance,
                                 const int64_t* neighbors_row_splits,
                                 const TFeat* out_features_gradient,
                                 bool normalize) {
    bool has_importance = inp_importance;
 
#define FN_PARAMETERS                                                          \
    filter_backprop, filter_dims, num_out, num_inp, inp_features,              \
            inp_importance, neighbors_index, neighbors_kernel_index,           \
            neighbors_importance, neighbors_row_splits, out_features_gradient, \
            normalize
 
#define CALL_TEMPLATE(HAS_IMPORTANCE)                                  \
    if (HAS_IMPORTANCE == has_importance)                              \
        _SparseConvBackropFilterCPU<TFeat, TOut, TIndex, TKernelIndex, \
                                    HAS_IMPORTANCE>(FN_PARAMETERS);
 
#define CALL_TEMPLATE2  \
    CALL_TEMPLATE(true) \
    CALL_TEMPLATE(false)
 
    CALL_TEMPLATE2
 
#undef CALL_TEMPLATE
#undef CALL_TEMPLATE2
 
#undef FN_PARAMETERS
}
 
}  // namespace impl
}  // namespace ml
}  // namespace open3d