ආදම්ප්රශස්තකරණය

40import math
41from typing import Dict, Any, Tuple, Optional
42
43import torch
44from labml import tracker
45from torch import nn
46
47from labml_nn.optimizers import GenericAdaptiveOptimizer, WeightDecay

50class Adam(GenericAdaptiveOptimizer):

58    def __init__(self, params,
59                 lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-16,
60                 weight_decay: WeightDecay = WeightDecay(),
61                 optimized_update: bool = True,
62                 defaults: Optional[Dict[str, Any]] = None):

76        defaults = {} if defaults is None else defaults
77        defaults.update(weight_decay.defaults())
78        super().__init__(params, defaults, lr, betas, eps)
79
80        self.weight_decay = weight_decay
81        self.optimized_update = optimized_update

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

93        state['step'] = 0

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

97        state['exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format)

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

109        beta1, beta2 = group['betas']

112        m, v = state['exp_avg'], state['exp_avg_sq']

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

119        v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
120
121        return m, v

123    def get_lr(self, state: Dict[str, any], group: Dict[str, any]):

131        return group['lr']

133    def adam_update(self, state: Dict[str, any], group: Dict[str, any], param: torch.nn.Parameter,
134                    m: torch.Tensor, v: torch.Tensor):

166        beta1, beta2 = group['betas']

168        bias_correction1 = 1 - beta1 ** state['step']

170        bias_correction2 = 1 - beta2 ** state['step']

173        lr = self.get_lr(state, group)

176        if self.optimized_update:

178            denominator = v.sqrt().add_(group['eps'])

180            step_size = lr * math.sqrt(bias_correction2) / bias_correction1

183            param.data.addcdiv_(m, denominator, value=-step_size)

185        else:

187            denominator = (v.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

189            step_size = lr / bias_correction1

192            param.data.addcdiv_(m, denominator, value=-step_size)

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

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

208        m, v = self.get_mv(state, group, grad)

211        state['step'] += 1

214        self.adam_update(state, group, param, m, v)