Package: autoclip
AutoClipPytorch and tensorflow implementations (and variations) of the AutoClip gradient smoothing procedure fromSeetharaman et al.Prem Seetharaman, Gordon Wichern, Bryan Pardo, Jonathan Le Roux. "AutoClip: Adaptive Gradient Clipping for Source Separation Networks." 2020 IEEE 30th International Workshop on Machine Learning for Signal Processing (MLSP). IEEE, 2020.AboutWhile training your model, AutoClip keeps a running history of all of your model's gradient magnitudes. Using these, the gradient clipper can adaptively clamp outlier gradient values before they reach the optimizer of your choice.While AutoClip is great as a preventative measure against exploding gradients, it also speeds up training time, and encourages the optimizer to find more optimal models. At an intuitive level, AutoClip compensates for the stochastic nature of training over batches, regularizing training effects.InstallationAutoClip is listed on pypi. To install AutoClip simply run the following commandpip install autoclipand theautoclippackage will be installed in your currently active environment.Torch APIBelow are some examples how to useautoclip's torch API.Clippers as Optimizer WrappersUsing the optimizer wrapping pattern is the recommended way to use AutoClip, andautoclip's torch API supports wrapping arbitrary pytorch optimizers. The wrapping pattern allows you to avoid changing your training code when you want to use an AutoClip clipper. This is especially useful if you do not own the training code for whatever reason. (Say for example you are using someone else's Trainer class, as is often the case with frameworks likehuggingface.)The following is an example of how to integrate AutoClip into your model training using this pattern:importtorchfromautoclip.torchimportQuantileClipmodel=torch.nn.Sequential(torch.nn.Linear(100,50),torch.nn.ReLU(),torch.nn.Linear(50,2))optimizer=torch.optim.AdamW(model.parameters())optimizer=QuantileClip.as_optimizer(optimizer=optimizer,quantile=0.9,history_length=1000)Now you can use the optimizer just like you would have before adding the clipper, and the clipping will be applied automatically.Raw AutoClip ClippersYou can still use the clipper manually if you would like. If this is the case, then you would create your clipper like this:importtorchfromautoclip.torchimportQuantileClipmodel=torch.nn.Sequential(torch.nn.Linear(100,50),torch.nn.ReLU(),torch.nn.Linear(50,2))clipper=QuantileClip(model.parameters(),quantile=0.9,history_length=1000)Then, to clip the model's gradients, simply run the clipper's.step()function during your training loop. Note that you should call the clipper'sstepbefore you call your optimizer'sstep. Calling it after would mean that your clipping will have no effect, since the model will have already been updated using the unclipped gradients. For example:forbatch_num,batchinenumerate(training_dataset):model_prediction=model(batch['data'])loss=loss_function(model_prediction,batch['targets'])loss.backward()clipper.step()# clipper comes before optimizeroptimizer.step()Global vs Local Clippingautoclip's torch clippers support two clipping modes. The first isglobal_clipping, which is the original AutoClip as described in Seetherman et al. The second is local or parameter-wise clipping. In this mode a history is kept for every parameter, and each is clipped according to its own history. By default, theautoclipclippers will use the parameter-wise clipping.
To use the global mode, simply pass the appropriate flag:clipper=QuantileClip(model.parameters(),quantile=0.9,history_length=1000,global_clipping=True)CheckpointingThe torch clippers also support checkpointing throughstate_dict()andload_state_dict(), just like torch models and optimizers. For example, if you want to checkpoint a clipper toclipper.pth:clipper=QuantileClip(model.parameters())torch.save(clipper.state_dict(),'clipper.pth')# Then laterclipper=QuantileClip(model.parameters())clipper.load_state_dict(torch.load('clipper.pth'))Keep in mind that just like a torch optimizer this will error if you give the clipper differently sized model parameters.While it is generally recommended to usestate_dicts instead (see thepytorch documentationon this subject for more info), you may also usetorch.saveandtorch.loaddirectly to pickle the entire clipper object.Tensorflowautoclip's tensorflow API does not currently have feature parity with the torch API (If you want to change this, feel free tocontribute).
As it is, the tensorflow API currently only supports the original AutoClip algorithm, and does not support checkpointing. Below is a short example:importtensorflowastffromautoclip.tfimportQuantileClipmodel=tf.keras.models.Sequential([tf.keras.layers.Dense(50),tf.keras.layers.ReLU(),tf.keras.layers.Dense(10),tf.keras.layers.ReLU(),tf.keras.layers.Dense(2,activation=tf.keras.activations.tanh),])model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001,gradient_transformers=[QuantileClip(quantile=0.9,history_length=1000)]),loss="mean_absolute_error",metrics=["accuracy"],)model.fit(train_data,train_targets)