I am wondering what I am doing wrong with register_hook since it does not seem to register a hook.
My goal is to modify grad matrix before weights are updated. For debugging purposes I just assigned zero matrix for a grad variable, but network trains perfectly.
So clearly hooks don’t work, but I cannot figure why. Here is my train function:
def train_cnn():
model = Net()
model.cuda()
criterion = torch.nn.CrossEntropyLoss()
model.train()
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)
for e in range(epochs):
agg_loss = 0
for data in trainloader:
x, y = data
x = x.cuda()
y = y.cuda()
outputs = model(x)
loss = criterion(outputs, y)
optimizer.zero_grad()
loss.backward()
with torch.no_grad():
hooks = []
for conv_layer in model.convos:
conv_layer.weight.retain_grad()
h = conv_layer.weight.register_hook(lambda grad: torch.zeros(size=grad))
hooks.append(h)
for fc_layer in model.linears:
fc_layer.weight.retain_grad()
h = fc_layer.weight.register_hook(lambda grad: torch.zeros(size=grad))
hooks.append(h)
optimizer.step()
for h in hooks:
h.remove()
Thank you!