Tell Pytorch's Autograd not to include a operation for gradient computation

1

I have a custom layer. Let the layer be called ‘Gaussian’

class Gaussian(nn.Module):
  def __init__():  
    super(Gaussian, self).__init__()
 
 #@torch.no_grad     
  def forward(self, x):
    x  = x + torch.rand(x.shape) 
    x[2:]  = x[2:] + x[1] # Some similar operations are also performed
    return x

cnn_model = nn.Sequential(nn.Conv2d(1, 32, 5), Gaussian(), nn.ReLU())

If a is the input, I want gradient of a to exclude operations that are present in the Gaussian module, but include the calculations in other layers of the neural network(nn.Conv2d etc).

In the end, my aim is to use the Gaussian module to perform calculations but that calculations should not be included in gradient computation.

I tried to do the following:

  1. Used the @torch.no_grad above the forward method of the Gaussian

  2. Using detach after every operation in the Gaussian module: x = x + torch.rand(x.shape) and similarly for other operations

Are the above methods correct?

2

Both approaches would break the computation graph and the previous layers wouldn’t get a valid gradient.
You could try to adapt this workflow (or write a custom autograd.Function and define the backward method manually).

1 Like
3

The modification helped me, thank you