def set_optimizer_momentum(optimizer: Optimizer, value: float, index: int = 0):
"""Set momentum of ``index`` 'th param group of optimizer to ``value``.
Args:
optimizer: PyTorch optimizer
value (float): new value of momentum
index (int, optional): integer index of optimizer's param groups,
default is 0
"""
betas = optimizer.param_groups[0].get("betas", None)
momentum = optimizer.param_groups[0].get("momentum", None)
if betas is not None:
_, beta = betas
optimizer.param_groups[index]["betas"] = (value, beta)
elif momentum is not None:
optimizer.param_groups[index]["momentum"] = value
def grad_step(*,
optimizer: Optimizer,
optimizer_wds: List[float] = 0,
grad_clip_fn: Callable = None):
"""
Makes a gradient step for a given optimizer
Args:
optimizer (Optimizer): the optimizer
optimizer_wds (List[float]): list of weight decay parameters
for each param group
grad_clip_fn (Callable): function for gradient clipping
"""
for group, wd in zip(optimizer.param_groups, optimizer_wds):
if wd > 0:
for param in group["params"]:
param.data = param.data.add(-wd * group["lr"], param.data)
if grad_clip_fn is not None:
grad_clip_fn(group["params"])
optimizer.step()