_conv_weight_size Class — pytorch Architecture

Source Code

aten/src/ATen/native/ConvUtils.h lines 292–307

template <typename T>
std::vector<T> _conv_weight_size(
    ArrayRef<T> input_size, ArrayRef<T> output_size,
    ArrayRef<T> padding, ArrayRef<T> output_padding, IntArrayRef stride, IntArrayRef dilation, int64_t groups
) {
  auto dim = input_size.size();
  std::vector<T> weight_size(dim);
  weight_size[0] = output_size[1];
  weight_size[1] = input_size[1] / groups;
  for (const auto d : c10::irange(2, dim)) {
    auto kernel = input_size[d] - (output_size[d] - 1) * stride[d - 2]
               + padding[d - 2] * 2 - output_padding[d - 2];
    weight_size[d] = (kernel - 1) / dilation[d - 2] + 1;
  }
  return weight_size;
}