如何在 PyTorch 中用零填充张量？

Question

A really simple question here but I can't seem to figure this one out. As part of preparation to finals in Deep learning, I'm trying to solve questions from previous exams. I need to write a similar Conv2d method which does stride and padding. My current code is:

class MyConv2d(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, kernel_size=(1,1), stride=1, padding=0):
        # Set input as fields
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.p = kernel_size[0]
        self.q = kernel_size[1]
        self.stride = stride
        self.padding = padding
        self.kern = nn.Parameter(torch.rand((out_channels, in_channels, kernel_size[0], kernel_size[1])))
        self.bias = nn.Parameter(torch.rand((out_channels,1)))
        
    def forward(self, X):
        X_clone = X.clone()
        h = int((X_clone.size(2) + (2 * self.padding) - self.p) / self.stride) + 1
        w = int((X_clone.size(3) + (2 * self.padding) - self.q) / self.stride) + 1
        result = torch.empty(X_clone.size(0), self.out_channels, h, w)
        

Test case:

batch_size = 3
H, W = 6,6
in_channels = 3
out_channels = 1
kernel_size = (2,2)
stride = 2
padding= 1
X = torch.rand(batch_size, in_channels, H, W)
conv = MyConv2d(in_channels=in_channels, out_channels=out_channels,kernel_size=kernel_size, stride=stride, padding=padding)
res = conv.forward(X)

The forward method should gave the same result as of Conv2d . The next step in forward method is padding with zeros but I can't seem to figure an easy way to pad X_clone with zeros. How it can be done?

EDIT : Forgot to mention that as part of the question, I'm not allowed to use any other methods under nn .

Answer 1

I can't figure out other fancy methods except creating a new tensor and adding the original one to it.

This padding function could be helpful:

def zero_padding(input_tensor, pad_size: int = 1):
  h, w = input_tensor.shape  # assuming no batch and channel dimension
  pad_tensor = torch.zeros([pad_size*2 + h, pad_size*2 + w])
  pad_tensor[pad_size:pad_size+h, pad_size:pad_size+w] = input_tensor
  return pad_tensor

You can expand this function to work with inputs containing batch and channel dimensions, arbitrary padding, etc.