crop.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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.

#ifndef LBANN_LAYER_CROP_HPP_INCLUDED
#define LBANN_LAYER_CROP_HPP_INCLUDED

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"

#include "lbann/layers/data_type_layer.hpp"
#include "lbann/proto/datatype_helpers.hpp"
#include "lbann/proto/layers.pb.h"
#include "lbann/utils/exception.hpp"
#include "lbann/utils/protobuf.hpp"

namespace lbann {

template <typename TensorDataType,
          data_layout T_layout = data_layout::DATA_PARALLEL,
          El::Device Dev = El::Device::CPU>
class crop_layer : public data_type_layer<TensorDataType>
{
  static_assert(T_layout == data_layout::DATA_PARALLEL,
                "crop layer only supports DATA_PARALLEL");
#ifdef LBANN_HAS_GPU_FP16
  using CompareType =
    typename std::conditional<std::is_same<TensorDataType, fp16>::value,
                              float,
                              TensorDataType>::type;
#else
  using CompareType = TensorDataType;
#endif
public:

  using AbsDistMatrixType = El::AbstractDistMatrix<TensorDataType>;


public:
  crop_layer(lbann_comm* comm, std::vector<int> dims)
    : data_type_layer<TensorDataType>(nullptr)
  {
    this->set_output_dims(dims);
    this->m_expected_num_parent_layers = 2;
  }

  crop_layer(const crop_layer& other)
    : data_type_layer<TensorDataType>(other),
      m_input_v(other.m_input_v ? other.m_input_v->Copy() : nullptr),
      m_output_v(other.m_output_v ? other.m_output_v->Copy() : nullptr),
      m_crop_pos_v(other.m_crop_pos_v ? other.m_crop_pos_v->Copy() : nullptr)
  {}
  crop_layer& operator=(const crop_layer& other)
  {
    data_type_layer<TensorDataType>::operator=(other);
    m_input_v.reset(other.m_input_v ? other.m_input_v->Copy() : nullptr);
    m_output_v.reset(other.m_output_v ? other.m_output_v->Copy() : nullptr);
    m_crop_pos_v.reset(other.m_crop_pos_v ? other.m_crop_pos_v->Copy()
                                          : nullptr);
    return *this;
  }

  crop_layer* copy() const override { return new crop_layer(*this); }


  template <typename ArchiveT>
  void serialize(ArchiveT& ar);


  std::string get_type() const override { return "crop"; }
  data_layout get_data_layout() const override { return T_layout; }
  El::Device get_device_allocation() const override { return Dev; }
  bool can_run_inplace() const override { return false; }
  int get_backprop_requirements() const override { return ERROR_SIGNALS; }

  void setup_data(size_t max_mini_batch_size) override
  {
    data_type_layer<TensorDataType>::setup_data(max_mini_batch_size);
    const auto& input = this->get_prev_activations();
    const auto& dist = input.DistData();
    m_input_v.reset(input.Construct(input.Grid(), input.Root()));
    m_output_v.reset(input.Construct(input.Grid(), input.Root()));

    m_crop_pos_v.reset(AbsDistMatrixType::Instantiate(
      *dist.grid,
      dist.root,
      El::STAR,
      dist.rowDist,
      (dist.blockWidth == 1 ? El::ELEMENT : El::BLOCK),
      El::Device::CPU));
  }

  void setup_dims() override
  {
    data_type_layer<TensorDataType>::setup_dims();
    std::stringstream err;

    // Make sure input tensors have valid dimensions
    const auto& input_dims = this->get_input_dims(0);
    const auto& loc_dims = this->get_input_dims(1);
    const auto& output_dims = this->get_output_dims();
    if (input_dims.size() != output_dims.size()) {
      err << get_type() << " layer \"" << this->get_name() << "\" "
          << "expects a crop input tensor with " << output_dims.size()
          << " dimensions, "
          << "but parent layer "
          << "\"" << this->get_parent_layers()[0]->get_name() << "\" "
          << "outputs a tensor with " << input_dims.size() << " dimensions";
      LBANN_ERROR(err.str());
    }
    if (loc_dims.size() != 1 || loc_dims[0] != (int)input_dims.size()) {
      err << get_type() << " layer \"" << this->get_name() << "\" "
          << "expects a 1D crop position tensor with " << output_dims.size()
          << " entries, "
          << "but parent layer "
          << "\"" << this->get_parent_layers()[1]->get_name() << "\" "
          << "outputs a tensor with dimensions ";
      for (size_t i = 0; i < loc_dims.size(); ++i) {
        err << (i > 0 ? " x " : "") << loc_dims[i];
      }
      LBANN_ERROR(err.str());
    }
  }

protected:
  void write_specific_proto(lbann_data::Layer& proto) const final;

  friend class cereal::access;
  crop_layer() : crop_layer(nullptr, {1}) {}

  void fp_compute() override
  {
    switch (this->get_input_dims().size()) {
    case 3:
      fp_compute_3d();
      break;
    default:
      fp_compute_nd();
    }
  }

  void bp_compute() override
  {
    switch (this->get_input_dims().size()) {
    case 3:
      bp_compute_3d();
      break;
    default:
      bp_compute_nd();
    }
  }

private:
  std::unique_ptr<AbsDistMatrixType> m_input_v;
  std::unique_ptr<AbsDistMatrixType> m_output_v;
  std::unique_ptr<AbsDistMatrixType> m_crop_pos_v;

  void fp_compute_nd()
  {

    // Input and output tensors
    const auto& input = this->get_prev_activations(0);
    auto& output = this->get_activations();

    // Tensor dimensions
    const auto& input_dims = this->get_input_dims(0);
    const auto& output_dims = this->get_output_dims();
    const El::Int num_dims = output_dims.size();
    const auto& local_width = input.LocalWidth();
    const auto& region_size = output_dims.back();

    // Get crop position
    m_crop_pos_v->Empty(false);
    m_crop_pos_v->AlignWith(input);
    const auto& input1 = this->get_prev_activations(1);
    if (m_crop_pos_v->DistData() == input1.DistData()) {
      El::LockedView(*m_crop_pos_v, input1);
    }
    else {
      El::Copy(input1, *m_crop_pos_v);
    }
    const auto& local_crop_pos = m_crop_pos_v->LockedMatrix();

    // Crop each local mini-batch sample
    // BVE_FIXME LBANN_OMP_PARALLEL_FOR
    for (El::Int local_col = 0; local_col < local_width; ++local_col) {
      const auto& col = input.GlobalCol(local_col);

      // Determine crop position
      std::vector<El::Int> crop_offsets;
      for (El::Int d = 0; d < num_dims; ++d) {
        const auto& pos = local_crop_pos(d, local_col);
        if (CompareType(pos) < CompareType(0.0) ||
            CompareType(pos) > CompareType(1.0)) {
          std::stringstream err;
          err << "crop position not in range [0,1] (pos=(";
          for (El::Int i = 0; i < local_crop_pos.Height(); ++i) {
            err << (i > 0 ? "," : "") << local_crop_pos(i, local_col);
          }
          err << "))";
          LBANN_ERROR(err.str());
        }
        const El::Int num_offsets = input_dims[d] - output_dims[d] + 1;
        crop_offsets.push_back(
          std::min(El::Int(static_cast<CompareType>(pos) * num_offsets),
                   num_offsets - 1));
      }

      // Copy contiguous regions from input tensor to output tensor
      std::vector<El::Int> output_pos(num_dims, 0);
      while (output_pos[0] < output_dims[0]) {

        // Copy region from input tensor to output tensor
        auto input_index = output_pos[0] + crop_offsets[0];
        auto output_index = output_pos[0];
        for (El::Int d = 1; d < num_dims; ++d) {
          input_index =
            (input_dims[d] * input_index + output_pos[d] + crop_offsets[d]);
          output_index = output_dims[d] * output_index + output_pos[d];
        }
        El::LockedView(*m_input_v,
                       input,
                       El::IR(input_index, input_index + region_size),
                       El::IR(col));
        El::View(*m_output_v,
                 output,
                 El::IR(output_index, output_index + region_size),
                 El::IR(col));
        El::Copy(*m_input_v, *m_output_v);

        // Move to next contiguous region
        output_pos.back() += region_size;
        for (El::Int d = num_dims - 1; d >= 1; --d) {
          if (output_pos[d] >= output_dims[d]) {
            output_pos[d] = 0;
            ++output_pos[d - 1];
          }
        }
      }
    }
  }

  void bp_compute_nd()
  {

    // Clear error signals
    El::Zero(this->get_error_signals(0));
    El::Zero(this->get_error_signals(1));

    // Input and gradient tensors
    const auto& gradient_wrt_output = this->get_prev_error_signals();
    auto& gradient_wrt_input = this->get_error_signals(0);
    const auto& local_crop_pos = m_crop_pos_v->LockedMatrix();

    // Tensor dimensions
    const auto& input_dims = this->get_input_dims(0);
    const auto& output_dims = this->get_output_dims();
    const El::Int num_dims = output_dims.size();
    const auto& local_width = gradient_wrt_input.LocalWidth();
    const auto& region_size = output_dims.back();

    // Populate error signal for each local mini-batch sample
    // BVE_FIXME LBANN_OMP_PARALLEL_FOR
    for (El::Int local_col = 0; local_col < local_width; ++local_col) {
      const auto& col = gradient_wrt_input.GlobalCol(local_col);

      // Determine crop position
      std::vector<El::Int> crop_offsets;
      for (El::Int d = 0; d < num_dims; ++d) {
        const auto& pos = local_crop_pos(d, local_col);
        if (CompareType(pos) < CompareType(0.0) ||
            CompareType(pos) > CompareType(1.0)) {
          std::stringstream err;
          err << "crop position not in range [0,1] (pos=(";
          for (El::Int i = 0; i < local_crop_pos.Height(); ++i) {
            err << (i > 0 ? "," : "") << local_crop_pos(i, local_col);
          }
          err << "))";
          LBANN_ERROR(err.str());
        }
        const El::Int num_offsets = input_dims[d] - output_dims[d] + 1;
        crop_offsets.push_back(
          std::min(El::Int(static_cast<CompareType>(pos) * num_offsets),
                   num_offsets - 1));
      }

      // Populate contiguous regions in gradient w.r.t. input tensor
      std::vector<El::Int> output_pos(num_dims, 0);
      while (output_pos[0] < output_dims[0]) {

        // Copy region
        auto input_index = output_pos[0] + crop_offsets[0];
        auto output_index = output_pos[0];
        for (El::Int d = 1; d < num_dims; ++d) {
          input_index =
            (input_dims[d] * input_index + output_pos[d] + crop_offsets[d]);
          output_index = output_dims[d] * output_index + output_pos[d];
        }
        El::LockedView(*m_output_v,
                       gradient_wrt_output,
                       El::IR(output_index, output_index + region_size),
                       El::IR(col));
        El::View(*m_input_v,
                 gradient_wrt_input,
                 El::IR(input_index, input_index + region_size),
                 El::IR(col));
        El::Copy(*m_output_v, *m_input_v);

        // Move to next contiguous region
        output_pos.back() += region_size;
        for (El::Int d = num_dims - 1; d >= 1; --d) {
          if (output_pos[d] >= output_dims[d]) {
            output_pos[d] = 0;
            ++output_pos[d - 1];
          }
        }
      }
    }
  }

  void fp_compute_3d();
  void bp_compute_3d();
};

template <typename T, data_layout L, El::Device D>
void crop_layer<T, L, D>::write_specific_proto(lbann_data::Layer& proto) const
{
  proto.set_datatype(proto::ProtoDataType<T>);
  auto* msg = proto.mutable_crop();
  protobuf::assign_to_repeated(*msg->mutable_dims(), this->get_output_dims());
}

#ifndef LBANN_CROP_LAYER_INSTANTIATE
#define PROTO_DEVICE(T, Device)                                                \
  extern template class crop_layer<T, data_layout::DATA_PARALLEL, Device>

#include "lbann/macros/instantiate_device.hpp"
#undef PROTO_DEVICE
#endif // LBANN_CROP_LAYER_INSTANTIATE

} // namespace lbann

#pragma GCC diagnostic pop
#endif // LBANN_LAYER_CROP_HPP_INCLUDED