ContinuousConvTranspose.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018-2021 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#pragma once

#include <tbb/parallel_for.h>

#include "open3d/ml/impl/continuous_conv/CoordinateTransformation.h"

namespace open3d {
namespace ml {
namespace impl {

template <class TFeat,
          class TOut,
          class TReal,
          class TIndex,
          InterpolationMode INTERPOLATION,
          CoordinateMapping MAPPING,
          bool ALIGN_CORNERS,
          bool INDIVIDUAL_EXTENT,
          bool ISOTROPIC_EXTENT,
          bool NORMALIZE>
void _CConvTransposeComputeFeaturesCPU(
        TOut* out_features,
        const std::vector<int>& filter_dims,
        const TFeat* filter,
        size_t num_out,
        const TReal* out_positions,
        const TFeat* out_importance,
        size_t num_inp,
        const TReal* inp_positions,
        const TFeat* inp_features,
        const TFeat* inp_neighbors_importance_sum,
        const int64_t* inp_neighbors_row_splits,
        size_t neighbors_index_size,
        const TIndex* neighbors_index,
        const TFeat* neighbors_importance,
        const int64_t* neighbors_row_splits,
        const TReal* extents,
        const TReal* offsets) {
    const bool NEIGHBORS_IMPORTANCE = inp_neighbors_importance_sum;
    const int VECSIZE = 32;
    typedef Eigen::Array<TReal, VECSIZE, 1> Vec_t;
    typedef InterpolationVec<TReal, VECSIZE, INTERPOLATION> InterpolationVec_t;
    InterpolationVec_t interpolation;

    const int in_channels = filter_dims[filter_dims.size() - 2];
    const int out_channels = filter_dims[filter_dims.size() - 1];

    int spatial_filter_size = 1;
    for (int i = 0; i < 3; ++i) spatial_filter_size *= filter_dims[i];
    Eigen::Array<int, 3, 1> filter_size_xyz(filter_dims[2], filter_dims[1],
                                            filter_dims[0]);

    memset(out_features, 0, sizeof(TOut) * num_out * out_channels);

    tbb::parallel_for(
            tbb::blocked_range<size_t>(0, num_out, 32),
            [&](const tbb::blocked_range<size_t>& r) {
                int range_length = r.end() - r.begin();

                Eigen::Matrix<TFeat, Eigen::Dynamic, Eigen::Dynamic> B(
                        in_channels * spatial_filter_size, range_length);
                B.setZero();

                typedef Eigen::Array<TFeat, VECSIZE, Eigen::Dynamic> Matrix;
                Matrix infeat(VECSIZE, in_channels);

                Eigen::Array<TReal, 3, 1> offsets_(offsets[0], offsets[1],
                                                   offsets[2]);

                Eigen::Array<TReal, VECSIZE, 3> inv_extents;
                if (INDIVIDUAL_EXTENT == false) {
                    if (ISOTROPIC_EXTENT) {
                        inv_extents = 1 / extents[0];
                    } else {
                        inv_extents.col(0) = 1 / extents[0];
                        inv_extents.col(1) = 1 / extents[1];
                        inv_extents.col(2) = 1 / extents[2];
                    }
                }

                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 =
                            (out_idx + 1 < num_out
                                     ? neighbors_row_splits[out_idx + 1]
                                     : neighbors_index_size);

                    typename InterpolationVec_t::Weight_t interp_weights;
                    typename InterpolationVec_t::Idx_t interp_indices;

                    int vec_valid_count = 0;
                    Vec_t x, y, z;

                    // set to zero to avoid problems with vectors with less than
                    // VECSIZE valid entries
                    x.setZero();
                    y.setZero();
                    z.setZero();
                    for (size_t n = neighbor_start; n < neighbor_end; ++n) {
                        const size_t inp_idx = neighbors_index[n];

                        const int i = vec_valid_count;
                        x(i) = out_positions[out_idx * 3 + 0] -
                               inp_positions[inp_idx * 3 + 0];
                        y(i) = out_positions[out_idx * 3 + 1] -
                               inp_positions[inp_idx * 3 + 1];
                        z(i) = out_positions[out_idx * 3 + 2] -
                               inp_positions[inp_idx * 3 + 2];

                        if (INDIVIDUAL_EXTENT) {
                            if (ISOTROPIC_EXTENT) {
                                inv_extents.row(i) = 1 / extents[inp_idx];
                            } else {
                                inv_extents(i, 0) =
                                        1 / extents[3 * inp_idx + 0];
                                inv_extents(i, 1) =
                                        1 / extents[3 * inp_idx + 1];
                                inv_extents(i, 2) =
                                        1 / extents[3 * inp_idx + 2];
                            }
                        }

                        TFeat n_importance = NEIGHBORS_IMPORTANCE
                                                     ? neighbors_importance[n]
                                                     : TFeat(1);
                        for (int ic = 0; ic < in_channels; ++ic)
                            infeat(i, ic) =
                                    inp_features[inp_idx * in_channels + ic] *
                                    n_importance;

                        if (NORMALIZE) {
                            TFeat normalizer(1);
                            if (NEIGHBORS_IMPORTANCE) {
                                if (inp_neighbors_importance_sum[inp_idx] !=
                                    TFeat(0))
                                    normalizer /= inp_neighbors_importance_sum
                                            [inp_idx];
                            } else {
                                size_t num_inp_neighbors;
                                const size_t inp_neighbor_start =
                                        inp_neighbors_row_splits[inp_idx];
                                const size_t inp_neighbor_end =
                                        inp_neighbors_row_splits[inp_idx + 1];
                                num_inp_neighbors =
                                        inp_neighbor_end - inp_neighbor_start;
                                if (num_inp_neighbors > 0)
                                    normalizer /= TFeat(num_inp_neighbors);
                            }
                            for (int ic = 0; ic < in_channels; ++ic)
                                infeat(i, ic) *= normalizer;
                        }

                        ++vec_valid_count;
                        if (vec_valid_count == VECSIZE ||
                            n + 1 == neighbor_end) {
                            ComputeFilterCoordinates<ALIGN_CORNERS, MAPPING>(
                                    x, y, z, filter_size_xyz, inv_extents,
                                    offsets_);
                            interpolation.Interpolate(
                                    interp_weights, interp_indices, x, y, z,
                                    filter_size_xyz, in_channels);
                            for (int k = 0; k < vec_valid_count; ++k) {
                                for (int j = 0; j < InterpolationVec_t::Size();
                                     ++j) {
                                    for (int ic = 0; ic < in_channels; ++ic)
                                        B(interp_indices(j, k) + ic, out_col) +=
                                                TFeat(interp_weights(j, k)) *
                                                infeat(k, ic);
                                }
                            }
                            vec_valid_count = 0;
                        }
                    }

                }  // out_idx

                Eigen::Map<const Eigen::Matrix<TFeat, Eigen::Dynamic,
                                               Eigen::Dynamic>>
                        A(filter, out_channels,
                          spatial_filter_size * in_channels);
                Eigen::Map<Eigen::Matrix<TOut, Eigen::Dynamic, Eigen::Dynamic>>
                        C(out_features + (r.begin() * out_channels),
                          out_channels, range_length);

                C = (A * B).template cast<TOut>();
                if (out_importance) {
                    for (int i = 0; i < range_length; ++i)
                        C.col(i) *= TOut(out_importance[r.begin() + i]);
                }
            });
}

template <class TFeat, class TOut, class TReal, class TIndex>
void CConvTransposeComputeFeaturesCPU(TOut* out_features,
                                      const std::vector<int>& filter_dims,
                                      const TFeat* filter,
                                      size_t num_out,
                                      const TReal* out_positions,
                                      const TFeat* out_importance,
                                      size_t num_inp,
                                      const TReal* inp_positions,
                                      const TFeat* inp_features,
                                      const TFeat* inp_neighbors_importance_sum,
                                      const int64_t* inp_neighbors_row_splits,
                                      size_t neighbors_index_size,
                                      const TIndex* neighbors_index,
                                      const TFeat* neighbors_importance,
                                      const int64_t* neighbors_row_splits,
                                      const TReal* extents,
                                      const TReal* offsets,
                                      InterpolationMode interpolation,
                                      CoordinateMapping coordinate_mapping,
                                      bool align_corners,
                                      bool individual_extent,
                                      bool isotropic_extent,
                                      bool normalize) {
#define FN_PARAMETERS                                                          \
    out_features, filter_dims, filter, num_out, out_positions, out_importance, \
            num_inp, inp_positions, inp_features,                              \
            inp_neighbors_importance_sum, inp_neighbors_row_splits,            \
            neighbors_index_size, neighbors_index, neighbors_importance,       \
            neighbors_row_splits, extents, offsets

#define CALL_TEMPLATE(INTERPOLATION, MAPPING, ALIGN_CORNERS,                \
                      INDIVIDUAL_EXTENT, ISOTROPIC_EXTENT, NORMALIZE)       \
    if (INTERPOLATION == interpolation && MAPPING == coordinate_mapping &&  \
        ALIGN_CORNERS == align_corners &&                                   \
        INDIVIDUAL_EXTENT == individual_extent &&                           \
        ISOTROPIC_EXTENT == isotropic_extent && NORMALIZE == normalize)     \
        _CConvTransposeComputeFeaturesCPU<TFeat, TOut, TReal, TIndex,       \
                                          INTERPOLATION, MAPPING,           \
                                          ALIGN_CORNERS, INDIVIDUAL_EXTENT, \
                                          ISOTROPIC_EXTENT, NORMALIZE>(     \
                FN_PARAMETERS);

#define CALL_TEMPLATE2(INTERPOLATION, MAPPING)                       \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, true, true)    \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, true, false)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, false, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, true, false, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, true, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, true, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, false, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, true, false, false, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, true, true)   \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, true, false)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, false, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, true, false, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, true, true)  \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, true, false) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, false, true) \
    CALL_TEMPLATE(INTERPOLATION, MAPPING, false, false, false, false)

#define CALL_TEMPLATE3(INTERPOLATION)                                     \
    CALL_TEMPLATE2(INTERPOLATION, CoordinateMapping::BALL_TO_CUBE_RADIAL) \
    CALL_TEMPLATE2(INTERPOLATION,                                         \
                   CoordinateMapping::BALL_TO_CUBE_VOLUME_PRESERVING)     \
    CALL_TEMPLATE2(INTERPOLATION, CoordinateMapping::IDENTITY)

#define CALL_TEMPLATE4                               \
    CALL_TEMPLATE3(InterpolationMode::LINEAR)        \
    CALL_TEMPLATE3(InterpolationMode::LINEAR_BORDER) \
    CALL_TEMPLATE3(InterpolationMode::NEAREST_NEIGHBOR)

    CALL_TEMPLATE4

#undef CALL_TEMPLATE
#undef CALL_TEMPLATE2
#undef CALL_TEMPLATE3
#undef CALL_TEMPLATE4

#undef FN_PARAMETERS
}

}  // namespace impl
}  // namespace ml
}  // namespace open3d