o ,iA@sddlmZmZmZddlZddlmZddlmZmZm Z m Z m Z m Z m Z mZmZmZmZmZddgZGdddeZd d e d ed ed e d e_d eedeedeedeedeededededededededefddZd eedeedeedeedeededededededededefddZe ed     d"d eedeedeedeedeed eededededededededefd!dZdS)#)ListOptionalTupleN)Tensor) _capturable_doc_default_to_fused_or_foreach_differentiable_doc_disable_dynamo_if_unsupported _foreach_doc!_get_capturable_supported_devices_get_scalar_dtype _maximize_doc_use_grad_for_differentiable _view_as_real OptimizerParamsTRproprpropcseZdZ   dddddddeded eeefd eeefd ed eed edeffddZfddZ ddZ e dddZ Z S)r{Gz?g?g333333?gư>2FN) capturableforeachmaximizedifferentiableparamslretas step_sizesrrrrc sd|ks td|d|dkrdkr|dks-ntd|dd|dt|||||||d} t|| dS) NgzInvalid learning rate: r?rzInvalid eta values: z, )rrr rrrr) ValueErrordictsuper__init__) selfrrrr rrrrdefaults __class__O/home/ubuntu/SoloSpeech/.venv/lib/python3.10/site-packages/torch/optim/rprop.pyr%s ( zRprop.__init__cst||jD]S}|dd|dd|dd|dd|dD]4}|j|g}t|dkr[t|ds[t |d}|drQtj |t |j d ntj |t d |d<q'q dS) NrrFrrrrstepdtypedevicer.) r$ __setstate__ param_groups setdefaultstategetlentorch is_tensorfloattensorr r/)r&r4grouppp_statestep_valr(r*r+r15s&         zRprop.__setstate__c Csd}|dD]}|jdurq|t|O}|||j} | jr$td|| |j|} t| dkrt|drBtjdt |j dntjdt d| d <tj |tj d | d <|j jrjt| t|d |d | d <n t| |d | d <|| d || d || d q|S)NFrz'Rprop does not support sparse gradientsrrr*r-r0r, memory_formatprevr step_size)gradr7 is_complexappend is_sparse RuntimeErrorr4r6zerosr r/ zeros_likepreserve_formatr. full_likecomplex) r&r;rgradsprevsr state_steps has_complexr<rCr4r*r*r+ _init_groupHs4       zRprop._init_groupcCs|d}|dur!t |}Wdn1swY|jD]@}g}g}g}g}g}|d\} } |d\} } |d} |d}|||||||}t|||||| | | | | ||d|d|dq$|S) zPerforms a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. Nrr rrrr) step_size_min step_size_maxetaminusetaplusrrrrrP) _cuda_graph_capture_health_checkr7 enable_gradr2rQr)r&closurelossr;rrMrNr rOrTrUrRrSrrrPr*r*r+r,nsH     z Rprop.step)rrrN)__name__ __module__ __qualname__rr9rboolrr%r1rQrr, __classcell__r*r*r(r+rs<     &a Implements the resilient backpropagation algorithm. .. math:: \begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \theta_0 \in \mathbf{R}^d \text{ (params)},f(\theta) \text{ (objective)}, \\ &\hspace{13mm} \eta_{+/-} \text{ (etaplus, etaminus)}, \Gamma_{max/min} \text{ (step sizes)} \\ &\textbf{initialize} : g^0_{prev} \leftarrow 0, \: \eta_0 \leftarrow \text{lr (learning rate)} \\ &\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} \textbf{for} \text{ } i = 0, 1, \ldots, d-1 \: \mathbf{do} \\ &\hspace{10mm} \textbf{if} \: g^i_{prev} g^i_t > 0 \\ &\hspace{15mm} \eta^i_t \leftarrow \mathrm{min}(\eta^i_{t-1} \eta_{+}, \Gamma_{max}) \\ &\hspace{10mm} \textbf{else if} \: g^i_{prev} g^i_t < 0 \\ &\hspace{15mm} \eta^i_t \leftarrow \mathrm{max}(\eta^i_{t-1} \eta_{-}, \Gamma_{min}) \\ &\hspace{15mm} g^i_t \leftarrow 0 \\ &\hspace{10mm} \textbf{else} \: \\ &\hspace{15mm} \eta^i_t \leftarrow \eta^i_{t-1} \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1}- \eta_t \mathrm{sign}(g_t) \\ &\hspace{5mm}g_{prev} \leftarrow g_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 the paper `A Direct Adaptive Method for Faster Backpropagation Learning: The RPROP Algorithm `_. a Args: params (iterable): iterable of parameters to optimize or dicts defining parameter groups lr (float, optional): learning rate (default: 1e-2) etas (Tuple[float, float], optional): pair of (etaminus, etaplus), that are multiplicative increase and decrease factors (default: (0.5, 1.2)) step_sizes (Tuple[float, float], optional): a pair of minimal and maximal allowed step sizes (default: (1e-6, 50)) z z rrMrNr rOrRrSrTrUrrrrPcCst|D]\} }|| }| s|n| }|| }|| }|| }tjs?| r?t}|jj|jjkr7|jj|vs?Jd|d|d7}t|r\t|}t|}t|}t|}| rh| |  }n| | }| r| t |d||| t |d||| t |dd|n|||d<|||d<d||d<|||||j tjd}| r| t ||d|nd|||<|j| |dd| |qdS)NIIf capturable=True, params and state_steps must be on supported devices: .rrr?value) enumerater7_utils is_compilingr r/typerD view_as_realmulclonesigncopy_wheregtlteqmul_clamp_rJaddcmul_)rrMrNr rOrRrSrTrUrrrrPiparamrCrArBr,capturable_supported_devicesrlr*r*r+_single_tensor_rpropsH       rxc s,t|dkrdS| rJdtjs.| r.ttfddt||Ds.Jddt|||||g} | D]\\}}}}}}|dj rXtj |tj ddd dd nt |d | rgt ||||t||}| rtt|t||| rt||}t|| r|D])}|t|d|||t|d|||t|dd |qn|D]}|||d<|||d<d ||d<qt|||D]}|||qt|}tt|D]}||t|||d||q~d d |D}tj|||ddq4s   z&_multi_tensor_rprop..r`rar!cpu)r/)alpharcSsg|]}|qSr*)rl)ryrCr*r*r+ ~sz'_multi_tensor_rprop..rbrc)r6r7rfrgr allzipr"_group_tensors_by_device_and_dtypevaluesis_cpu _foreach_add_r:r _foreach_mul _foreach_neg__foreach_copy__foreach_sign_rmrnrorprq _foreach_mul_rslistrange_foreach_addcmul_)rrMrNr rOrRrSrTrUrrrrPgrouped_tensorsgrouped_params grouped_grads grouped_prevsgrouped_step_sizesgrouped_state_steps_signsrlrBru grad_signsr*rzr+_multi_tensor_rprops|              r)single_tensor_fnFrc Cstjstdd|Dstd|durt||dd\}}|r*tjr*td|r4tjs4t}nt }||||||| | | | |||| d dS) zpFunctional API that performs rprop algorithm computation. See :class:`~torch.optim.Rprop` for details. css|] }t|tjVqdSrZ) isinstancer7r)rytr*r*r+r{s  zrprop..zPAPI has changed, `state_steps` argument must contain a list of singleton tensorsNF) use_fusedz6torch.jit.script not supported with foreach optimizers)rRrSrTrUrrrrP) r7rfrgrrGrjit is_scriptingrrx)rrMrNr rOrrrrrPrRrSrTrUrfuncr*r*r+rs<  )NFFFF)typingrrrr7r optimizerrrr r r r r rrrrr__all__r__doc__r9r^rxrrr*r*r*r+s 8 #    7      D      l