src/operator/correlation.cc - master - githubmirror/incubator-mxnet

src/operator/correlation.cc
210 строк11 KB
mozga
[master][clang-format] Re-format cc. .h. .cu files; cond. (#20704)
19 ноя 2021, 11:27
19 ноя 2021, 11:271add250
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file correlation.cc
 * \brief correlation op
 * \author Xu Dong
 */
#include "./correlation-inl.h"
#include "./mshadow_op.h"

namespace mshadow {
template <typename Dtype>
void AddPad(const Tensor<cpu, 4, Dtype>& original, const Tensor<cpu, 4, Dtype>& out, int pad_size) {
  for (index_t nbatch = 0; nbatch < original.size(0); nbatch++)
    for (index_t channel = 0; channel < original.size(1); channel++)
      for (index_t h = 0; h < original.size(2); h++)
        for (index_t w = 0; w < original.size(3); w++)
          out[nbatch][h + pad_size][w + pad_size][channel] = original[nbatch][channel][h][w];
}
template <typename Dtype>
inline void CorrelationForward(const Tensor<cpu, 4, Dtype>& out,
                               const Tensor<cpu, 4, Dtype>& data1,
                               const Tensor<cpu, 4, Dtype>& data2,
                               const Tensor<cpu, 4, Dtype>& tmp1,
                               const Tensor<cpu, 4, Dtype>& tmp2,
                               int top_channels_,
                               int top_height_,
                               int top_width_,
                               int pad_size_,
                               bool is_multiply,
                               int max_displacement_,
                               int kernel_size_,
                               int neighborhood_grid_radius_,
                               int neighborhood_grid_width_,
                               int kernel_radius_,
                               int stride1_,
                               int stride2_) {
  const index_t bnum  = data1.size(0);
  const int bchannels = data1.size(1);
  const int sumelems  = kernel_size_ * kernel_size_ * bchannels;
  AddPad<Dtype>(data1, tmp1, pad_size_);
  index_t top_channels_unsigned_ = static_cast<index_t>(top_channels_);
  AddPad<Dtype>(data2, tmp2, pad_size_);
  for (index_t i = 0; i < static_cast<index_t>(top_height_); i++)
    for (index_t j = 0; j < static_cast<index_t>(top_width_); j++)
      for (index_t nbatch = 0; nbatch < bnum; nbatch++) {
        int x1 = j * stride1_ + max_displacement_;
        int y1 = i * stride1_ + max_displacement_;
        for (index_t top_channel = 0; top_channel < top_channels_unsigned_; top_channel++) {
          int s2o = (top_channel % neighborhood_grid_width_ - neighborhood_grid_radius_) * stride2_;
          int s2p = (top_channel / neighborhood_grid_width_ - neighborhood_grid_radius_) * stride2_;
          int x2  = x1 + s2o;
          int y2  = y1 + s2p;
          for (index_t h = 0; h < static_cast<index_t>(kernel_size_); h++)
            for (index_t w = 0; w < static_cast<index_t>(kernel_size_); w++)
              for (index_t channel = 0; channel < static_cast<index_t>(bchannels); channel++) {
                if (is_multiply == true)
                  out[nbatch][top_channel][i][j] +=
                      tmp1[nbatch][y1 + h][x1 + w][channel] * tmp2[nbatch][y2 + h][x2 + w][channel];
                else
                  out[nbatch][top_channel][i][j] += std::abs(tmp1[nbatch][y1 + h][x1 + w][channel] -
                                                             tmp2[nbatch][y2 + h][x2 + w][channel]);
              }
          out[nbatch][top_channel][i][j] /= sumelems;
        }
      }
}
template <typename Dtype>
inline void CorrelationBackward(const Tensor<cpu, 4, Dtype>& out_grad,
                                const Tensor<cpu, 4, Dtype>& in_grad1,
                                const Tensor<cpu, 4, Dtype>& in_grad2,
                                const Tensor<cpu, 4, Dtype>& tmp1,
                                const Tensor<cpu, 4, Dtype>& tmp2,
                                int top_channels_,
                                int top_height_,
                                int top_width_,
                                int pad_size_,
                                bool is_multiply,
                                int max_displacement_,
                                int kernel_size_,
                                int neighborhood_grid_radius_,
                                int neighborhood_grid_width_,
                                int kernel_radius_,
                                int stride1_,
                                int stride2_,
                                int num,
                                int channels,
                                int height,
                                int width) {
  const float sumelems = kernel_size_ * kernel_size_ * channels;
  for (index_t i = 0; i < static_cast<index_t>(top_height_); i++)
    for (index_t j = 0; j < static_cast<index_t>(top_width_); j++)
      for (index_t nbatch = 0; nbatch < static_cast<index_t>(num); nbatch++) {
        int x1 = j * stride1_ + max_displacement_;
        int y1 = i * stride1_ + max_displacement_;
        for (int top_channel = 0; top_channel < top_channels_; top_channel++) {
          int s2o = (top_channel % neighborhood_grid_width_ - neighborhood_grid_radius_) * stride2_;
          int s2p = (top_channel / neighborhood_grid_width_ - neighborhood_grid_radius_) * stride2_;
          int x2  = x1 + s2o;
          int y2  = y1 + s2p;
          for (int h = 0; h < kernel_size_; h++)
            for (int w = 0; w < kernel_size_; w++)
              for (int channel = 0; channel < channels; channel++) {
                if (is_multiply == true) {
                  if ((y1 + h - pad_size_ >= 0) && (x1 + w - pad_size_ >= 0) &&
                      (y1 + h < height + pad_size_) && (x1 + w < width + pad_size_)) {
                    in_grad1[nbatch][channel][y1 + h - pad_size_][x1 + w - pad_size_] +=
                        out_grad[nbatch][top_channel][i][j] *
                        tmp2[nbatch][y2 + h][x2 + w][channel] / sumelems;
                  }
                  if ((y2 + h - pad_size_ >= 0) && (x2 + w - pad_size_ >= 0) &&
                      (y2 + h < height + pad_size_) && (x2 + w < width + pad_size_)) {
                    in_grad2[nbatch][channel][y2 + h - pad_size_][x2 + w - pad_size_] +=
                        out_grad[nbatch][top_channel][i][j] *
                        tmp1[nbatch][y1 + h][x1 + w][channel] / sumelems;
                  }
                } else {
                  if ((y1 + h - pad_size_ >= 0) && (x1 + w - pad_size_ >= 0) &&
                      (y1 + h < height + pad_size_) && (x1 + w < width + pad_size_)) {
                    Dtype sign = (tmp1[nbatch][y1 + h][x1 + w][channel] >=
                                  tmp2[nbatch][y2 + h][x2 + w][channel]) ?
                                     Dtype(1.0) :
                                     Dtype(-1.0);
                    in_grad1[nbatch][channel][y1 + h - pad_size_][x1 + w - pad_size_] +=
                        out_grad[nbatch][top_channel][i][j] * sign / sumelems;
                  }
                  if ((y2 + h - pad_size_ >= 0) && (x2 + w - pad_size_ >= 0) &&
                      (y2 + h < height + pad_size_) && (x2 + w < width + pad_size_)) {
                    Dtype sign = (tmp1[nbatch][y1 + h][x1 + w][channel] >=
                                  tmp2[nbatch][y2 + h][x2 + w][channel]) ?
                                     Dtype(-1.0) :
                                     Dtype(1.0);
                    in_grad2[nbatch][channel][y2 + h - pad_size_][x2 + w - pad_size_] +=
                        out_grad[nbatch][top_channel][i][j] * sign / sumelems;
                  }
                }
              }
        }
      }
}
}  // namespace mshadow
namespace mxnet {
namespace op {
template <>
Operator* CreateOp<cpu>(CorrelationParam param, int dtype) {
  Operator* op = nullptr;
  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, { op = new CorrelationOp<cpu, DType>(param); });
  return op;
}
Operator* CorrelationProp::CreateOperatorEx(Context ctx,
                                            mxnet::ShapeVector* in_shape,
                                            std::vector<int>* in_type) const {
  DO_BIND_DISPATCH(CreateOp, param_, in_type->at(0));
}
DMLC_REGISTER_PARAMETER(CorrelationParam);
MXNET_REGISTER_OP_PROPERTY(Correlation, CorrelationProp)
    .add_argument("data1", "NDArray-or-Symbol", "Input data1 to the correlation.")
    .add_argument("data2", "NDArray-or-Symbol", "Input data2 to the correlation.")
    .add_arguments(CorrelationParam::__FIELDS__())
    .describe(R"code(Applies correlation to inputs.

The correlation layer performs multiplicative patch comparisons between two feature maps.

Given two multi-channel feature maps :math:`f_{1}, f_{2}`, with :math:`w`, :math:`h`, and :math:`c` being their width, height, and number of channels,
the correlation layer lets the network compare each patch from :math:`f_{1}` with each patch from :math:`f_{2}`.

For now we consider only a single comparison of two patches. The 'correlation' of two patches centered at :math:`x_{1}` in the first map and
:math:`x_{2}` in the second map is then defined as:

.. math::

   c(x_{1}, x_{2}) = \sum_{o \in [-k,k] \times [-k,k]} <f_{1}(x_{1} + o), f_{2}(x_{2} + o)>

for a square patch of size :math:`K:=2k+1`.

Note that the equation above is identical to one step of a convolution in neural networks, but instead of convolving data with a filter, it convolves data with other
data. For this reason, it has no training weights.

Computing :math:`c(x_{1}, x_{2})` involves :math:`c * K^{2}` multiplications. Comparing all patch combinations involves :math:`w^{2}*h^{2}` such computations.

Given a maximum displacement :math:`d`, for each location :math:`x_{1}` it computes correlations :math:`c(x_{1}, x_{2})` only in a neighborhood of size :math:`D:=2d+1`,
by limiting the range of :math:`x_{2}`. We use strides :math:`s_{1}, s_{2}`, to quantize :math:`x_{1}` globally and to quantize :math:`x_{2}` within the neighborhood
centered around :math:`x_{1}`.

The final output is defined by the following expression:

.. math::
  out[n, q, i, j] = c(x_{i, j}, x_{q})

where :math:`i` and :math:`j` enumerate spatial locations in :math:`f_{1}`, and :math:`q` denotes the :math:`q^{th}` neighborhood of :math:`x_{i,j}`.
)code" ADD_FILELINE);
}  // namespace op
}  // namespace mxnet