Как использовать слой с градиентом, но без регулировки веса? - PullRequest
Новая
       107

Как использовать слой с градиентом, но без регулировки веса?

2 голосов
/ 20 июня 2020

Можно ли отметить часть прямого прохода только для обратного распространения градиента, но не для корректировки весов?

В следующем примере кода у меня есть Module, который использует только один слой (один набор параметры), но он используется дважды на шаге вперед. Во время оптимизации я бы ожидал, что веса также будут скорректированы дважды. Если я хочу отрегулировать веса только для одного из применений слоя, что я могу сделать? 

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()

Следующее не будет работать, так как с torch.no_grad() градиент не распространяется в обратном направлении: 

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

Я не могу просто исключить параметры из оптимизации, так как их нужно оптимизировать в первой части (т.е. out1 = self.linear(x)). По той же причине я также не могу установить для этих параметров скорость обучения 0.

Что еще я могу сделать для этого?

1 Ответ

1 голос
/ 21 июня 2020

Один из способов сделать это - использовать requires_grad_ для временного отключения градиентов в параметрах слоя: 

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

Это все еще позволяет градиентам течь через активации; он просто не дает им достичь параметров.

Вы также можете рассмотреть возможность манипулирования тензорами веса вручную и вызова .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
...