Hello!
I’m trying to implement a MaxPooling2D layer from scratch without using any modules from toch.nn
.
def forward(self, x, debug=False):
# In shape
batch, in_features, h_in, w_in = x.shape
# Expected out dims
# If the img has been padded the term + 2 * padding should not be considered
h_out = int(((h_in - self.kernel) / self.stride)) + 1
w_out = int(((w_in - self.kernel) / self.stride)) + 1
# Preparing output
out = torch.zeros(size=(batch, in_features, h_out * w_out))
# Getting all the windows (basically what fold does)
offset = 0
for i in range(0, h_in, self.stride):
for j in range(0, w_in, self.stride):
if j + self.kernel > w_in or i + self.kernel > h_in:
continue
window_elems = x[:, :, i:i + self.kernel, j:j + self.kernel].long()
max_elems = torch.amax(window_elems, in_features)
elems = torch.amax(max_elems, 2)
out[:, :, offset] = elems
offset += 1
return out.reshape(batch, in_features, h_out, w_out)
The problem that I’m currently having is with the out tensor that not only is considered a leaf node but also it has no grad_fn (and that means I can not call the backward from the model). Is there any way that I can create the out vector while “linking it” with the input vector in the autograd graph?
Thanks in advance!