cpu_max_unpool_channels_last Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/MaxUnpoolKernel.cpp lines 101–162

template <typename scalar_t>
void cpu_max_unpool_channels_last(
    Tensor& output_,
    const Tensor& input,
    const Tensor& indices) {
  TORCH_CHECK(input.ndimension() == 4,
              "max_unpool2d with channels last format supports tensors with 4 dims");
  auto memory_format = at::MemoryFormat::ChannelsLast;
  auto output = output_.contiguous(memory_format);

  auto input_data = input.const_data_ptr<scalar_t>();
  auto indices_data = indices.const_data_ptr<int64_t>();
  auto output_data = output.data_ptr<scalar_t>();

  int64_t nbatch = input.size(0);
  int64_t channels = input.size(1);
  int64_t input_height = input.size(2);
  int64_t input_width = input.size(3);
  int64_t output_height = output.size(2);
  int64_t output_width = output.size(3);
  int64_t input_image_size = input_height * input_width;
  int64_t output_image_size = output_height * output_width;

  std::optional<int64_t> optional_error_index;

  // parallel on dim N, H, W
  at::parallel_for(0, nbatch * input_image_size, 0, [&](int64_t begin, int64_t end) {
    int64_t n = 0;
    int64_t ip = 0;
    data_index_init(begin, n, nbatch, ip, input_image_size);

    for (const auto i : c10::irange(begin, end)) {
      const scalar_t* input_ptr = input_data + i * channels;
      const int64_t* indices_ptr = indices_data + i * channels;
      scalar_t* output_ptr = output_data + n * output_image_size * channels;

      // can't do scatter on avx2 (only available on avx512)
      for (const auto c : c10::irange(channels)) {
        int64_t maxp = indices_ptr[c];
        if (maxp < 0 || maxp >= output_image_size) {
          optional_error_index = maxp;
          std::atomic_thread_fence(std::memory_order_release);
        } else {
          output_ptr[maxp * channels + c] = input_ptr[c];
        }
      }

      // move on to next input index
      data_index_step(n, nbatch, ip, input_image_size);
    }
  });

  if (optional_error_index) {
    TORCH_CHECK(false, "Found an invalid max index: ", optional_error_index.value(),
        " (output volumes are of size ", output_height,
        "x", output_width, ")");
  }

  if (!output_.is_contiguous(memory_format)) {
    output_.copy_(output);
  }
}