is_3d Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/MaxUnpoolKernel.cpp lines 16–99

template <typename scalar_t, bool is_3d = false>
void cpu_max_unpool(
    Tensor& output_,
    const Tensor& input,
    const Tensor& indices) {
  auto output = output_.contiguous();

  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>();

  // NB: input tensor dimensions:
  // MaxUnpool2d:
  //    dim = 3: CHW
  //    dim = 4: NCHW
  // MaxUnpool3d:
  //    dim = 4: CDHW
  //    dim = 5: NCDHW

  int64_t numel = input.numel();
  int64_t ndim = input.ndimension();

  // treat batch size and channels as one dimension
  // and the feature map as another dimension
  int64_t channels = 0;
  [[maybe_unused]] int64_t output_depth = 0;
  [[maybe_unused]] int64_t output_height = 0;
  [[maybe_unused]] int64_t output_width = 0;
  if constexpr (is_3d) {
    TORCH_CHECK(ndim == 4 || ndim == 5, "MaxUnpool3d: expect input to be 4d or 5d tensor.");
    channels = ndim == 4 ? input.size(0) : input.size(0) * input.size(1);
    output_depth = output.size(-3);
    output_height = output.size(-2);
    output_width = output.size(-1);
  } else {
    TORCH_CHECK(ndim == 3 || ndim == 4, "MaxUnpool2d: expect input to be 3d or 4d tensor.");
    channels = ndim == 3 ? input.size(0) : input.size(0) * input.size(1);
    output_depth = 1;
    output_height = output.size(-2);
    output_width = output.size(-1);
  }
  int64_t input_image_size = numel / channels;
  int64_t output_image_size = output.numel() / channels;

  std::optional<int64_t> optional_error_index;

  // parallel on dim N, C, D, H, W: [channels, input_image_size]
  at::parallel_for(0, numel, 0, [&](int64_t begin, int64_t end) {
    int64_t c = 0;
    int64_t ip = 0;
    data_index_init(begin, c, channels, ip, input_image_size);

    for (const auto i : c10::irange(begin, end)) {
      scalar_t* output_ptr = output_data + c * output_image_size;

      int64_t maxp = indices_data[i];
      if (maxp < 0 || maxp >= output_image_size) {
        optional_error_index = maxp;
        std::atomic_thread_fence(std::memory_order_release);
      } else {
        output_ptr[maxp] = input_data[i];
      }

      // move on to next input index
      data_index_step(c, channels, ip, input_image_size);
    }
  });

  if (optional_error_index) {
    if constexpr (is_3d) {
      TORCH_CHECK(false, "Found an invalid max index: ", optional_error_index.value(),
          " (output volumes are of size ", output_depth,
          "x", output_height, "x", output_width, ")");
    } else {
      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()) {
    output_.copy_(output);
  }
}