Conv2D source code

blade · July 1, 2022, 8:02pm

I was trying to look at the implementation of the 2D convolution on PyTorch. I found Conv2D class, which is using F.conv2, but I can’t get to the source code for the actual implementation, except for these 2 lines:

@overload
def conv2d(input: Tensor, weight: Tensor, bias: Optional[Tensor]=None, stride: Union[_int, _size]=1, padding: Union[_int, _size]=0, dilation: Union[_int, _size]=1, groups: _int=1) -> Tensor: ...
@overload
def conv2d(input: Tensor, weight: Tensor, bias: Optional[Tensor]=None, stride: Union[_int, _size]=1, padding: str="valid", dilation: Union[_int, _size]=1, groups: _int=1) -> Tensor: ...

Would appreciate if someone can direct me to the right place. Thanks!

Afekst · July 1, 2022, 8:21pm

Of course because PyTorch is open-source you can see it, but when using it within a python code you actually run a pre-compiled C++ code so it isn’t available.
The actual implementation can be found on GitHub

github.com

pytorch/pytorch/blob/e1cfac42b2b3ea96d3783c66c8d2e44244a40524/aten/src/ATen/native/Convolution.cpp#L799


      
            std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 1, "conv1d");
            Tensor output;
            if (at::isComplexType(input_.scalar_type())) {
              output = complex_convolution(input, weight, bias, stride, padding, dilation, {0}, groups);
            } else {
              output = at::convolution(input, weight, bias, stride, padding, dilation, false, {0}, groups);
            }
            return is_batched ? output : output.squeeze(0);
          }
          
          
at::Tensor conv2d(
              const Tensor& input_, const Tensor& weight, const c10::optional<Tensor>& bias_opt,
              IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
            // See [Note: hacky wrapper removal for optional tensor]
            c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
            const Tensor& bias = *bias_maybe_owned;
          
          
  Tensor input;
            bool is_batched;
            std::tie(input, is_batched) = batchify(input_, /*num_spatial_dims=*/ 2, "conv2d");
            Tensor output;

This is true for a lot of “basic” things in PyTorch.