[Solved] AssertionError: Trying to set 2 values for theta but there are 1 parameter groups

I had an issue that stumped me for a while so I want to share the solution. I kept getting this error when I was trying to pass a new thing that I wanted my optimizer to track for me.

@log_args(to_return=True, but_as=Optimizer.__init__)
def NAG(params, lr, mom=0., wd=0., decouple_wd=True):
    "A `Optimizer` for SGD with `lr` and `mom` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    if mom != 0: cbs.append(average_grad)
    cbs.append(sgd_step if mom==0 else momentum_step)
    return Optimizer(params, cbs, lr=lr, mom=mom, wd=wd, theta=params)

NAG(L(learn.model.parameters()), 0.1, mom=0.9)

This is the output that I would get when trying to run this. (I also tried really breaking down the problem and get rid of the learner I had defined, but that didn’t end up being where the problem was.
Output stack trace:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-74-36e23e3549e8> in <module>
----> 1 NAG(L(learn.model.parameters()), 0.1, mom=0.9)

~/Environment_personal/development/fastcore/fastcore/utils.py in _f(*args, **kwargs)
    467         log_dict = {**func_args.arguments, **{f'{k} (not in signature)':v for k,v in xtra_kwargs.items()}}
    468         log = {f'{f.__qualname__}.{k}':v for k,v in log_dict.items() if k not in but}
--> 469         inst = f(*args, **kwargs) if to_return else args[0]
    470         init_args = getattr(inst, 'init_args', {})
    471         init_args.update(log)

<ipython-input-72-7296e5482c82> in NAG(params, lr, mom, wd, decouple_wd)
      5     if mom != 0: cbs.append(average_grad)
      6     cbs.append(sgd_step if mom==0 else momentum_step)
----> 7     return Optimizer(params, cbs, lr=lr, mom=mom, wd=wd, theta=params)

~/Environment_personal/development/fastcore/fastcore/utils.py in _f(*args, **kwargs)
    471         init_args.update(log)
    472         setattr(inst, 'init_args', init_args)
--> 473         return inst if to_return else f(*args, **kwargs)
    474     return _f
    475 

~/Environment_personal/development/fastai/fastai/optimizer.py in __init__(self, params, cbs, train_bn, **defaults)
     71         self.param_lists = L(L(p) for p in params) if isinstance(params[0], (L,list)) else L([params])
     72         self.hypers = L({} for _ in range_of(self.param_lists))
---> 73         self.set_hypers(**defaults)
     74         self.frozen_idx = 0
     75 

~/Environment_personal/development/fastai/fastai/optimizer.py in set_hypers(self, **kwargs)
     32 
     33     def unfreeze(self): self.freeze_to(0)
---> 34     def set_hypers(self, **kwargs): L(kwargs.items()).starmap(self.set_hyper)
     35     def _set_hyper(self, k, v):
     36         for v_,h in zip(v, self.hypers): h[k] = v_

~/Environment_personal/development/fastcore/fastcore/foundation.py in starmap(self, f, *args, **kwargs)
    417     def cycle(self): return cycle(self)
    418     def map_dict(self, f=noop, *args, **kwargs): return {k:f(k, *args,**kwargs) for k in self}
--> 419     def starmap(self, f, *args, **kwargs): return self._new(itertools.starmap(partial(f,*args,**kwargs), self))
    420     def zip(self, cycled=False): return self._new((zip_cycle if cycled else zip)(*self))
    421     def zipwith(self, *rest, cycled=False): return self._new([self, *rest]).zip(cycled=cycled)

~/Environment_personal/development/fastcore/fastcore/foundation.py in _new(self, items, *args, **kwargs)
    340     @property
    341     def _xtra(self): return None
--> 342     def _new(self, items, *args, **kwargs): return type(self)(items, *args, use_list=None, **kwargs)
    343     def __getitem__(self, idx): return self._get(idx) if is_indexer(idx) else L(self._get(idx), use_list=None)
    344     def copy(self): return self._new(self.items.copy())

~/Environment_personal/development/fastcore/fastcore/foundation.py in __call__(cls, x, *args, **kwargs)
     49             return x
     50 
---> 51         res = super().__call__(*((x,) + args), **kwargs)
     52         res._newchk = 0
     53         return res

~/Environment_personal/development/fastcore/fastcore/foundation.py in __init__(self, items, use_list, match, *rest)
    331         if items is None: items = []
    332         if (use_list is not None) or not _is_array(items):
--> 333             items = list(items) if use_list else _listify(items)
    334         if match is not None:
    335             if is_coll(match): match = len(match)

~/Environment_personal/development/fastcore/fastcore/foundation.py in _listify(o)
    244     if isinstance(o, list): return o
    245     if isinstance(o, str) or _is_array(o): return [o]
--> 246     if is_iter(o): return list(o)
    247     return [o]
    248 

~/Environment_personal/development/fastai/fastai/optimizer.py in set_hyper(self, k, v)
     42         v = L(v, use_list=None)
     43         if len(v)==1: v = v*len(self.param_lists)
---> 44         assert len(v) == len(self.hypers), f"Trying to set {len(v)} values for {k} but there are {len(self.param_lists)} parameter groups."
     45         self._set_hyper(k, v)
     46 

AssertionError: Trying to set 2 values for theta but there are 1 parameter groups.

It turned out that in order to do this, I needed to wrap the thing I wanted to track in another layer of list. This is my working code:

@log_args(to_return=True, but_as=Optimizer.__init__)
def NAG(params, lr, mom=0., wd=0., decouple_wd=True):
    "A `Optimizer` for SGD with `lr` and `mom` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    if mom != 0: cbs.append(average_grad)
    cbs.append(sgd_step if mom==0 else momentum_step)
    return Optimizer(params, cbs, lr=lr, mom=mom, wd=wd, theta=[params])
1 Like