[docs]classAdamax(Optimizer):def__init__(self,params,lr=2e-3,betas=(0.9,0.999),eps=1e-8,weight_decay=0,foreach:Optional[bool]=None,*,maximize:bool=False,differentiable:bool=False,):ifnot0.0<=lr:raiseValueError(f"Invalid learning rate: {lr}")ifnot0.0<=eps:raiseValueError(f"Invalid epsilon value: {eps}")ifnot0.0<=betas[0]<1.0:raiseValueError(f"Invalid beta parameter at index 0: {betas[0]}")ifnot0.0<=betas[1]<1.0:raiseValueError(f"Invalid beta parameter at index 1: {betas[1]}")ifnot0.0<=weight_decay:raiseValueError(f"Invalid weight_decay value: {weight_decay}")defaults=dict(lr=lr,betas=betas,eps=eps,weight_decay=weight_decay,foreach=foreach,maximize=maximize,differentiable=differentiable,)super().__init__(params,defaults)def__setstate__(self,state):super().__setstate__(state)forgroupinself.param_groups:group.setdefault("foreach",None)group.setdefault("maximize",False)group.setdefault("differentiable",False)state_values=list(self.state.values())step_is_tensor=(len(state_values)!=0)andtorch.is_tensor(state_values[0]["step"])ifnotstep_is_tensor:forsinstate_values:s["step"]=torch.tensor(float(s["step"]))def_init_group(self,group,params_with_grad,grads,exp_avgs,exp_infs,state_steps):forpingroup["params"]:ifp.gradisNone:continueparams_with_grad.append(p)ifp.grad.is_sparse:raiseRuntimeError("Adamax does not support sparse gradients")grads.append(p.grad)state=self.state[p]# State initializationiflen(state)==0:state["step"]=torch.tensor(0.0)state["exp_avg"]=torch.zeros_like(p,memory_format=torch.preserve_format)state["exp_inf"]=torch.zeros_like(p,memory_format=torch.preserve_format)exp_avgs.append(state["exp_avg"])exp_infs.append(state["exp_inf"])state_steps.append(state["step"])
[docs]@_use_grad_for_differentiabledefstep(self,closure=None):"""Performs a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. """loss=NoneifclosureisnotNone:withtorch.enable_grad():loss=closure()forgroupinself.param_groups:params_with_grad=[]grads=[]exp_avgs=[]exp_infs=[]state_steps=[]beta1,beta2=group["betas"]eps=group["eps"]lr=group["lr"]weight_decay=group["weight_decay"]foreach=group["foreach"]maximize=group["maximize"]differentiable=group["differentiable"]self._init_group(group,params_with_grad,grads,exp_avgs,exp_infs,state_steps)adamax(params_with_grad,grads,exp_avgs,exp_infs,state_steps,eps=eps,beta1=beta1,beta2=beta2,lr=lr,weight_decay=weight_decay,foreach=foreach,maximize=maximize,differentiable=differentiable,)returnloss
Adamax.__doc__=r"""Implements Adamax algorithm (a variant of Adam based on infinity norm). .. math:: \begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)}, \: \lambda \text{ (weight decay)}, \\ &\hspace{13mm} \epsilon \text{ (epsilon)} \\ &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, u_0 \leftarrow 0 \text{ ( infinity norm)} \\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ &\hspace{5mm}if \: \lambda \neq 0 \\ &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ &\hspace{5mm}u_t \leftarrow \mathrm{max}(\beta_2 u_{t-1}, |g_{t}|+\epsilon) \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \frac{\gamma m_t}{(1-\beta^t_1) u_t} \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned} For further details regarding the algorithm we refer to `Adam: A Method for Stochastic Optimization`_. """+fr""" Args: params (iterable): iterable of parameters to optimize or dicts defining parameter groups lr (float, optional): learning rate (default: 2e-3) betas (Tuple[float, float], optional): coefficients used for computing running averages of gradient and its square eps (float, optional): term added to the denominator to improve numerical stability (default: 1e-8) weight_decay (float, optional): weight decay (L2 penalty) (default: 0){_foreach_doc}{_maximize_doc}{_differentiable_doc} .. _Adam\: A Method for Stochastic Optimization: https://arxiv.org/abs/1412.6980 """defadamax(params:List[Tensor],grads:List[Tensor],exp_avgs:List[Tensor],exp_infs:List[Tensor],state_steps:List[Tensor],# kwonly args with defaults are not supported by functions compiled with torchscript issue #70627# setting this as kwarg for now as functional API is compiled by torch/distributed/optimforeach:Optional[bool]=None,maximize:bool=False,differentiable:bool=False,*,eps:float,beta1:float,beta2:float,lr:float,weight_decay:float,):r"""Functional API that performs adamax algorithm computation. See :class:`~torch.optim.Adamax` for details. """ifnotall(isinstance(t,torch.Tensor)fortinstate_steps):raiseRuntimeError("API has changed, `state_steps` argument must contain a list of singleton tensors")ifforeachisNone:_,foreach=_default_to_fused_or_foreach(params,differentiable,use_fused=False)ifforeachandtorch.jit.is_scripting():raiseRuntimeError("torch.jit.script not supported with foreach optimizers")ifforeachandnottorch.jit.is_scripting():func=_multi_tensor_adamaxelse:func=_single_tensor_adamaxfunc(params,grads,exp_avgs,exp_infs,state_steps,eps=eps,beta1=beta1,beta2=beta2,lr=lr,weight_decay=weight_decay,maximize=maximize,differentiable=differentiable,)def_single_tensor_adamax(params:List[Tensor],grads:List[Tensor],exp_avgs:List[Tensor],exp_infs:List[Tensor],state_steps:List[Tensor],*,eps:float,beta1:float,beta2:float,lr:float,weight_decay:float,maximize:bool,differentiable:bool,):fori,paraminenumerate(params):grad=grads[i]grad=gradifnotmaximizeelse-gradexp_avg=exp_avgs[i]exp_inf=exp_infs[i]step_t=state_steps[i]# update stepstep_t+=1ifweight_decay!=0:grad=grad.add(param,alpha=weight_decay)iftorch.is_complex(param):param=torch.view_as_real(param)grad=torch.view_as_real(grad)exp_avg=torch.view_as_real(exp_avg)exp_inf=torch.view_as_real(exp_inf)# Update biased first moment estimate.exp_avg.lerp_(grad,1-beta1)# Update the exponentially weighted infinity norm.norm_buf=torch.cat([exp_inf.mul_(beta2).unsqueeze(0),grad.abs().add_(eps).unsqueeze_(0)],0)ifnotdifferentiable:torch.amax(norm_buf,0,keepdim=False,out=exp_inf)else:exp_inf.copy_(torch.amax(norm_buf,0,keepdim=False))bias_correction=1-beta1**_get_value(step_t)clr=lr/bias_correctionparam.addcdiv_(exp_avg,exp_inf,value=-clr)def_multi_tensor_adamax(params:List[Tensor],grads:List[Tensor],exp_avgs:List[Tensor],exp_infs:List[Tensor],state_steps:List[Tensor],*,beta1:float,beta2:float,lr:float,weight_decay:float,eps:float,maximize:bool,differentiable:bool,):assertnotdifferentiable,"_foreach ops don't support autograd"iflen(params)==0:returngrouped_tensors=Optimizer._group_tensors_by_device_and_dtype([params,grads,exp_avgs,exp_infs,state_steps])for((grouped_params,grouped_grads,grouped_exp_avgs,grouped_exp_infs,grouped_state_steps),_)ingrouped_tensors.values():ifmaximize:grouped_grads=torch._foreach_neg(grouped_grads)grouped_params=[torch.view_as_real(x)iftorch.is_complex(x)elsexforxingrouped_params]grouped_grads=[torch.view_as_real(x)iftorch.is_complex(x)elsexforxingrouped_grads]grouped_exp_avgs=[torch.view_as_real(x)iftorch.is_complex(x)elsexforxingrouped_exp_avgs]grouped_exp_infs=[torch.view_as_real(x)iftorch.is_complex(x)elsexforxingrouped_exp_infs]# Update stepstorch._foreach_add_(grouped_state_steps,1)ifweight_decay!=0:ifmaximize:# Re-use the intermediate memory (device_grads) already allocated for maximizetorch._foreach_add_(grouped_grads,grouped_params,alpha=weight_decay)else:grouped_grads=torch._foreach_add(grouped_grads,grouped_params,alpha=weight_decay)# Update biased first moment estimate.torch._foreach_lerp_(grouped_exp_avgs,grouped_grads,1-beta1)# Update the exponentially weighted infinity norm.torch._foreach_mul_(grouped_exp_infs,beta2)forexp_inf,gradinzip(grouped_exp_infs,grouped_grads):norm_buf=torch.cat([exp_inf.unsqueeze(0),grad.abs().add_(eps).unsqueeze_(0)],0)torch.max(norm_buf,0,keepdim=False,out=(exp_inf,exp_inf.new().long()))bias_corrections=[1-beta1**_get_value(step)forstepingrouped_state_steps]clr=_stack_if_compiling([-1*(lr/bias_correction)forbias_correctioninbias_corrections])torch._foreach_addcdiv_(grouped_params,grouped_exp_avgs,grouped_exp_infs,clr)
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.