This is a tiny wrapper adding a max norm constraint as discussed in Hinton et al. (2012) and Srivastava et al. (2014) for PyTorch (I've used it on PyTorch 1.13.1) that wraps it around a PyTorch optimizer, in this case torch.optim.SGD
.
You can use MNSGD
as a drop-in replacement for SGD
, but you do need to pass it params as an iterable of dicts, and you need to pass 'max_norm'
as keyword for each parameter group to which you want to apply the max-norm constraint, like this:
model = NeuralNetwork(dropout_p = 0.5).to(device)
loss_fn = nn.CrossEntropyLoss()
model = NeuralNetwork().to(device)
optimizer = MNSGD(
[
{'params': model.hidden.parameters(), 'max_norm': 2 },
{'params': model.output.parameters() }
],
lr=0.5,
momentum=0.5
)
See PyTorch's doc for torch.optim.Optimizer
for more. The result of the above is that parameters matrices contained in model.hidden
will get clipped, and parameter matrix in model.output
won't be. Note that this is something of a hack that relies on torch.optim.Optimizer
accepting keyword arguments that don't match what it internally uses.
Here is an example of use that includes a test of correctness.
The max-norm constraint is essentially an additional normalization constraint applied after each optimization steps; the optimizer checks for each unit if its incoming weight vector tf.keras.constraint.MaxNorm
), but I couldn't find any out-of-the-box implementation for PyTorch.
Add other optimizers and support for non-iterable params
argument to the optimizer constructor. Run more tests, including with CUDA.