How to use custom convolution module instead of torch.nn.conv2d?

1

I am trying to use the example given below to construct a custom convolution module which will be used in the network instead of the torch.nn.conv2d.

inp = torch.randn(1, 3, 10, 12)
w = torch.randn(2, 3, 4, 5, requires_grad = True)
inp_unf = torch.nn.functional.unfold(inp, (4, 5))
out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))

Link to the example: [https://github.com/pytorch/pytorch/blob/e8536c08a16b533fe0a9d645dd4255513f9f4fdd/torch/nn/modules/fold.py]

How do I I go about this? Do I just make a class and write a forward function? If I have to make a class, what would be the base module?
Thank you.

2

You could just derive nn.Module like when creating a model. In the __init__ function you should create your parameters and the forward will perform your computation.
Since you used just pure PyTorch methods, autograd will be able to create the backward call.

1 Like