I am doing a final project for a class and I need to do a proof of concept for an idea we have. Implementing the whole idea is out of scope for the course but something we need to do for the proof of concept is to be able to use different feature maps in the backward pass than the forward pass. Right now I only care about doing this for a Conv2D layer. I’ve messed with different stuff like trying to mess with the gradient graph (obviously didn’t work) and I think my best option is to define an autograd function. For the forward pass I know I can just use the functional API, but for the backward is there anything in pytorch that would allow me to easily get the gradient of a conv2d layer so i can easily compute the conv2d gradient in the backward function? Also how can i use an autograd function as a layer in a Sequential model? What would I have to wrap it in? The idea is that in the forward pass we compute the feature maps but throw them out and in the backward pass we need to recompute the feature maps to use for gradient calculations, but instead of recomputing them with the normal weights, we would recompute them with an approximated version of the weights.
nn.grad had the gradient functions I wanted. My only remaining question is how can I turn a custom autograd function into a layer I can input into a Sequential model?
You can define a custom autograd.Function as described here and use it in a custom nn.Module by falling it inside the forward method of the module.