Package: gradinit
PyTorch implementation of the GradInitPytorch-based library to initialize any neural network with the gradient descent and your dataset.
This implementation is the simple way to use the method described in the paperChen Zhu et al. GradInit: Learning to Initialize Neural Networks for Stable and Efficient TrainingGradinit may be used as the alternative to the warmup mechanism.
It is more useful for deep or unusual architectures.Installationpip install gradinitorpip install --upgrade git+https://github.com/johngull/gradinit.gitDifferences from the original paperThis implementation uses simplified loss in the case of the small gradients.
in this case,gradinituses the user's loss instead of the complex 2-stage procedure described in the paper.UsageTo make gradinit you need to wrap your model with theGradInitWrapperand then you can use the usual training loop with some differences:recalculate loss with theGradInitWrapper.grad_lossfunctionclip scales with theGradInitWrapper.clamp_scalesfunctionHere is a simplified example of the full gradinit process:model:torch.nn.Module=...withGradInitWrapper(model)asginit:#it is important to create optimizer after wraping your modeloptimizer=torch.optim.Adam(model.parameters())norm=1# use 2 for the SGD optimizermodel.train()forx,yindata_loader:pred=model(x)loss=criterion(pred,y)# recalculate gradinit loss based on your lossloss=ginit.grad_loss(loss,norm)optimizer.zero_grad()loss.backward()optimizer.step()# clip minimum values for the init scalesginit.clamp_scales()# on exit of with statement model is recovered to its normal way# here you should create your main optimizer and start trainingoptimizer=torch.optim.Adam(model.parameters())...Alternatively to thewith-notation
you may usedetachor delete the wrapper object after finishing the initialization process.ginit=GradInitWrapper(model)# full gradinit loop...ginit.detach()# or del ginit# here you should create your main optimizer and start trainingoptimizer=torch.optim.Adam(model.parameters())...From our experience gradinit works worse in the mixed-precision mode.
And we recommend running gradinit in the full-precision mode and then starting the main training loop in mixed-precision.But if you really want to try, gradinit supports torch mixed-precision.
In such a case gradinit need to use your scaler object.
Here is an example of how to use gradinit with the torch amp.importtorch.cuda.ampmodel:torch.nn.Module=...scaler:torch.cuda.amp.GradScaler=...withGradInitWrapper(model)asginit:# it is important to create optimizer after wraping your modeloptimizer=torch.optim.Adam(model.parameters())norm=1# use 2 for the SGD optimizermodel.train()forx,yindata_loader:withautocast():pred=model(x)loss=criterion(pred,y)# recalculate gradinit loss based on your lossloss=ginit.grad_loss(loss,norm,scaler=scaler)optimizer.zero_grad()scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()# clip minimum values for the init scalesginit.clamp_scales()# on exit of with statement model is recovered to its normal way# here you should create your main optimizer and start trainingoptimizer=torch.optim.Adam(model.parameters())...ExperimentsFor the experiment results seethis page