U ʹj26 @sddlZddlmZddlmZmZmZmZmZmZm Z ddl m Z m Z ddgZ GdddeZd d e d ed ed d e_de ee ee e eee eeeeeeed ddZe ee ee e eeeeeeeed ddZe ee ee e eeeeeeeed ddZdS)N)Tensor) Optimizerrequired_use_grad_for_differentiable_default_to_fused_or_foreach_differentiable_doc _foreach_doc _maximize_doc)ListOptionalSGDsgdcsbeZdZeddddfddddeeeedfddZfddZd d Ze d d d Z Z S)r rFN)maximizeforeachdifferentiablec s|tk r|dkrtd||dkr4td||dkrJtd|t|||||||| d} |r~|dksv|dkr~tdt|| dS)NgzInvalid learning rate: zInvalid momentum value: zInvalid weight_decay value: )lrmomentum dampening weight_decaynesterovrrrrz8Nesterov momentum requires a momentum and zero dampening)r ValueErrordictsuper__init__) selfparamsrrrrrrrrdefaults __class__N/var/www/html/bwave/python/venv/lib/python3.8/site-packages/torch/optim/sgd.pyr s z SGD.__init__csLt||jD]4}|dd|dd|dd|ddqdS)NrFrrr)r __setstate__ param_groups setdefault)rstategrouprr r!r"s      zSGD.__setstate__cCsnd}|dD]\}|jdk r ||||j|jjr 1 \\ &\hspace{15mm} \textbf{b}_t \leftarrow \mu \textbf{b}_{t-1} + (1-\tau) g_t \\ &\hspace{10mm}\textbf{else} \\ &\hspace{15mm} \textbf{b}_t \leftarrow g_t \\ &\hspace{10mm}\textbf{if} \: \textit{nesterov} \\ &\hspace{15mm} g_t \leftarrow g_{t} + \mu \textbf{b}_t \\ &\hspace{10mm}\textbf{else} \\[-1.ex] &\hspace{15mm} g_t \leftarrow \textbf{b}_t \\ &\hspace{5mm}\textbf{if} \: \textit{maximize} \\ &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} + \gamma g_t \\[-1.ex] &\hspace{5mm}\textbf{else} \\[-1.ex] &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma g_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned} Nesterov momentum is based on the formula from `On the importance of initialization and momentum in deep learning`__. a Args: params (iterable): iterable of parameters to optimize or dicts defining parameter groups lr (float): learning rate momentum (float, optional): momentum factor (default: 0) weight_decay (float, optional): weight decay (L2 penalty) (default: 0) dampening (float, optional): dampening for momentum (default: 0) nesterov (bool, optional): enables Nesterov momentum (default: False) z z a Example: >>> # xdoctest: +SKIP >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> optimizer.zero_grad() >>> loss_fn(model(input), target).backward() >>> optimizer.step() __ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf .. note:: The implementation of SGD with Momentum/Nesterov subtly differs from Sutskever et. al. and implementations in some other frameworks. Considering the specific case of Momentum, the update can be written as .. math:: \begin{aligned} v_{t+1} & = \mu * v_{t} + g_{t+1}, \\ p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, \end{aligned} where :math:`p`, :math:`g`, :math:`v` and :math:`\mu` denote the parameters, gradient, velocity, and momentum respectively. This is in contrast to Sutskever et. al. and other frameworks which employ an update of the form .. math:: \begin{aligned} v_{t+1} & = \mu * v_{t} + \text{lr} * g_{t+1}, \\ p_{t+1} & = p_{t} - v_{t+1}. \end{aligned} The Nesterov version is analogously modified. Moreover, the initial value of the momentum buffer is set to the gradient value at the first step. This is in contrast to some other frameworks that initialize it to all zeros. ) rr,r-r.rrrrrrrc Csx|dkr*tjs&t|ddd\} }nd}|r@tjr@td|rTtjsTt} nt} | |||||||| || d dS)zlFunctional API that performs SGD algorithm computation. See :class:`~torch.optim.SGD` for details. NF)rZ use_fusedz6torch.jit.script not supported with foreach optimizers)rrrrrr.r)r1ZjitZ is_scriptingr RuntimeError_multi_tensor_sgd_single_tensor_sgd) rr,r-r.rrrrrrr_funcr r r!rs( ) rr,r-rrrrrrr.cCst|D]\} } |s|| n||  } |dkr<| j| |d} |dkr|| } | dkrlt| } | || <n| |j| d|d|r| j| |d} n| } | j| | dqdS)Nralphar) enumerateaddr1clonedetachmul_add_)rr,r-rrrrrrr.iparamZd_pbufr r r!r<s  r<) rgradsr-rrrrrrr.cCst|dkrdStj|||gdd} | D]\\} } } }tdd| D}|r\t| } |dkr|rztj| | |dntj| | |d} |dkrg}d}t t| D](}| |dkrd}qq| | |q|rt ||tj|| d|dn|g}t t| D]j}| |dkrBt | | }| |<|||<n$| |}||j| |d|d| |q|rtj| ||dn|} |stj| | | dq,t t| D]}| |j| || dqq,dS) NrT)Z with_indicescss|] }|jVqdS)N)r*).0r(r r r! sz$_multi_tensor_sgd..r?Fr)lenrZ"_group_tensors_by_device_and_dtypevaluesanyr1Z _foreach_negZ _foreach_add_Z _foreach_addranger)Z _foreach_mul_rCrDrErF)rrJr-rrrrrrr.Zgrouped_tensorsZ device_paramsZ device_gradsZdevice_momentum_buffer_listindicesZdevice_has_sparse_gradZbufsZall_states_with_momentum_bufferrGrIr r r!r; sJ     r;)NN)r1rZ optimizerrrrrrr r typingr r __all__r __doc__r8floatrr<r;r r r r!sh $V"   ._  .  #