Using PyTorch Optimizers

I am trying to use a trick that I saw in discord to show Nesterov Accelerated Gradient.

Here is the trick I saw: nesterov_opt = compose(torch.optim.SGD, OptimWrapper)

But when I tried running it with this code:

#hide
model = MySimpleModel()
model_tmp = MySimpleModel()
loss_func = mse #F.mse_loss
πœ‚ = 0.1
partial_learner = partialler(Learner,dls_normal, model, loss_func, cbs=[OptimizerInsight])
learn = partial_learner(partialler(nesterov_opt, momentum=0.75, nesterov=True), lr=πœ‚) #<<<<<<<<<<<<<<<<<

and then run learn.fit

I get this error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-156-3371a2bf699b> in <module>
----> 1 learn.fit(10)

~/Environment_personal/development/fastcore/fastcore/logargs.py in _f(*args, **kwargs)
     54         init_args.update(log)
     55         setattr(inst, 'init_args', init_args)
---> 56         return inst if to_return else f(*args, **kwargs)
     57     return _f

~/Environment_personal/development/fastai/fastai/learner.py in fit(self, n_epoch, lr, wd, cbs, reset_opt)
    200     def fit(self, n_epoch, lr=None, wd=None, cbs=None, reset_opt=False):
    201         with self.added_cbs(cbs):
--> 202             if reset_opt or not self.opt: self.create_opt()
    203             if wd is None: wd = self.wd
    204             if wd is not None: self.opt.set_hypers(wd=wd)

~/Environment_personal/development/fastai/fastai/learner.py in create_opt(self)
    139     def _bn_bias_state(self, with_bias): return norm_bias_params(self.model, with_bias).map(self.opt.state)
    140     def create_opt(self):
--> 141         self.opt = self.opt_func(self.splitter(self.model), lr=self.lr)
    142         if not self.wd_bn_bias:
    143             for p in self._bn_bias_state(True ): p['do_wd'] = False

~/Environment_personal/development/fastcore/fastcore/utils.py in _inner(x, *args, **kwargs)
    399     if order is not None: funcs = funcs.sorted(order)
    400     def _inner(x, *args, **kwargs):
--> 401         for f in L(funcs): x = f(x, *args, **kwargs)
    402         return x
    403     return _inner

TypeError: __init__() got an unexpected keyword argument 'momentum'

So I dug into it a little bit and it is passing all of the keywords into SGD and then into OptimWrapper. So it doesn’t matter if I pass momentum or mom. This other version does work for me though that Kushajveer Singh proposed:

def nesterov_opt(params, **kwargs): return OptimWrapper(torch.optim.SGD(params, **kwargs))

Just wondering if I messed up compose or if that won’t quite work as written.

1 Like