How to use a layer with gradient but without weight adjustment?

user_2230045 0
Viewed 256

Is it possible to mark part of the forward pass to only backpropagate the gradient but not to adjust weights?

In the following example code I have a Module that uses only one layer (one set of parameters) but it is used twice in the forward step. During the optimization I would expect the weights to be adjusted twice as well. If I want to only adjust the weights for one of the layer usages, what can I do?

import torch
    
class ExampleModel(torch.nn.Module):
    
    def __init__(self, dim) -> None:
        super(ExampleModel, self).__init__()
        self.linear = torch.nn.Linear(dim, dim)
    
    def forward(self, x):
        out1 = self.linear(x)  # backprop gradients and adjust weights here
        out2 = self.linear(out1)  # only backprop gradients here
        return out2
    
    
# Random input output data for this example
N, D = 64, 100
x = torch.randn(N, D)
y = torch.randn(N, D)
    
model = ExampleModel(D)
    
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.Adam(model.parameters())
    
y_pred = model(x)
loss = criterion(y_pred, y)
    
optimizer.zero_grad()
loss.backward()
optimizer.step()

The following will not work since with torch.no_grad() no gradient at all is backpropagated:

def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    with torch.no_grad():
        out2 = self.linear(out1)  # only backprop gradients here
    return out2

I can not simply exclude the parameters from the optimization since they need to be optimized in the first part (i.e. out1 = self.linear(x)). For the same reason I can also not set a learning rate of 0 for these parameters.

What else can I do to achieve this?

backpropagation pytorch optimization
1 Answers

One way to do it is to use requires_grad_ to temporarily disable gradients on the layer's parameters:

def forward(self, x):
    out1 = self.linear(x)  # backprop gradients and adjust weights here
    self.linear.requires_grad_(False)
    out2 = self.linear(out1)  # only backprop gradients here
    self.linear.requires_grad_(True)
    return out2

This still lets gradients flow through the activations; it merely stops them from reaching the parameters.

You could also consider manipulating the weight tensors manually and calling .detach():

import torch.nn.functional as F
def forward(self, x):
    out1 = self.linear(x)
    out2 = F.linear(out1, self.linear.weight.detach(), self.linear.bias.detach())
    return out2
user_2485220
Related