Hello all, I am trying to train a Resnet50 model on imagenet. My requirement is that I want to only train a subset of weights using a mask. I have successfully done this with general resnet models on cifar10 using the following sequence of code:
optimizer.zero_grad() loss.backward() for name, p in model.named_parameters(): if 'weight' in name: p.grad.data = p.grad.data * masks[name] optimizer.step()
masks is a dictionary containing layerwise masks with values of 1s and 0s.
Now, currently I am trying to incorporate the same idea for distributed imagenet training. I am using nVidia’s Resnet repo which utilizes AMP. I tried with including the following codeblock in nVidia’s implementation in this function in line 209 here:
scaler.scale(loss).backward() for name, p in model_and_loss.named_parameters(): if 'weight' in name: p.grad.data = p.grad.data * masks[name]
However, with the above code, still all of the weights are getting trained after gradient step. The optimizer step is done with (I think) here:
optimizer_step = ((i + 1) % batch_size_multiplier) == 0 loss = step(input, target, optimizer_step=optimizer_step) if ema is not None: ema(model_and_loss, epoch*steps_per_epoch+i)
I am unsure, how to achieve this in this nvidia implementation, also unsure how the gradient step is being taken.
My goal was to zero-out a subset of the gradients before taking the optimizer.step(), however with AMP it does not seem to be straightforward, specially in nvidia’s repo.
My question is, is there a different way to accomplish this in the training code here.
- Moreover, during AMP gradient update, is the same gradients from parameters.grad are used or am I missing something?
Thanks a lot for reading this. Please let me know if I need to clarify something here, I hope to be quick to response. I appreciate the help.