ADABeliefප්රශස්තකරණය

36from typing import Dict, Any
37
38import torch
39from torch import nn
40
41from labml_nn.optimizers import WeightDecay
42from labml_nn.optimizers.radam import RAdam

45class AdaBelief(RAdam):

52    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
53                 weight_decay: WeightDecay = WeightDecay(), amsgrad=False,
54                 degenerate_to_sgd=True,
55                 rectify=True, defaults=None):

73        defaults = {} if defaults is None else defaults
74        super().__init__(params, lr, betas, eps, weight_decay, amsgrad, degenerate_to_sgd, defaults)
75        self.rectify = rectify

77    def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):

85        state['step'] = 0

87        state['exp_avg'] = torch.zeros_like(param, memory_format=torch.preserve_format)

89        state['exp_avg_var'] = torch.zeros_like(param, memory_format=torch.preserve_format)

93        if group['amsgrad']:

95            state['max_exp_avg_var'] = torch.zeros_like(param, memory_format=torch.preserve_format)

97    def get_ms(self, state: Dict[str, Any], group: Dict[str, Any], grad: torch.Tensor):

107        beta1, beta2 = group['betas']

110        m, s = state['exp_avg'], state['exp_avg_var']

114        m.mul_(beta1).add_(grad, alpha=1 - beta1)

116        grad_residual = grad - m

119        s.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1 - beta2)

122        if group['amsgrad']:

124            s_max = state['max_exp_avg_var']

126            torch.maximum(s_max, s, out=s_max)
127
128            return m, s_max
129        else:

131            return m, s

133    def step_param(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor, param: torch.nn.Parameter):

144        grad = self.weight_decay(param, grad, group)

147        m, s = self.get_ms(state, group, grad)

150        state['step'] += 1
151
152        if not self.rectify:

155            self.adam_update(state, group, param, m, s + group['eps'])
156        else:

159            self.r_adam_update(state, group, param, m, s + group['eps'])