Inheritance diagram for fastreid.solver.optim.lamb.Lamb:
Public Member Functions | |
| __init__ (self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0, adam=False) | |
| step (self, closure=None) | |
Public Attributes | |
| adam | |
Detailed Description
Implements Lamb algorithm.
It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
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)
adam (bool, optional): always use trust ratio = 1, which turns this into
Adam. Useful for comparison purposes.
.. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
https://arxiv.org/abs/1904.00962
Constructor & Destructor Documentation
◆ __init__()
| fastreid.solver.optim.lamb.Lamb.__init__ | ( | self, | |
| params, | |||
| lr = 1e-3, | |||
| betas = (0.9, 0.999), | |||
| eps = 1e-6, | |||
| weight_decay = 0, | |||
| adam = False ) |
Definition at line 44 of file lamb.py.
45 weight_decay=0, adam=False):
46 if not 0.0 <= lr:
47 raise ValueError("Invalid learning rate: {}".format(lr))
48 if not 0.0 <= eps:
49 raise ValueError("Invalid epsilon value: {}".format(eps))
50 if not 0.0 <= betas[0] < 1.0:
51 raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
52 if not 0.0 <= betas[1] < 1.0:
53 raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
54 defaults = dict(lr=lr, betas=betas, eps=eps,
55 weight_decay=weight_decay)
56 self.adam = adam
57 super(Lamb, self).__init__(params, defaults)
58
Member Function Documentation
◆ step()
| fastreid.solver.optim.lamb.Lamb.step | ( | self, | |
| closure = None ) |
Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
Definition at line 59 of file lamb.py.
59 def step(self, closure=None):
60 """Performs a single optimization step.
61 Arguments:
62 closure (callable, optional): A closure that reevaluates the model
63 and returns the loss.
64 """
65 loss = None
66 if closure is not None:
67 loss = closure()
68
69 for group in self.param_groups:
70 for p in group['params']:
71 if p.grad is None or group['freeze']:
72 continue
73 grad = p.grad.data
74 if grad.is_sparse:
75 raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instad.')
76
77 state = self.state[p]
78
79 # State initialization
80 if len(state) == 0:
81 state['step'] = 0
82 # Exponential moving average of gradient values
83 state['exp_avg'] = torch.zeros_like(p.data)
84 # Exponential moving average of squared gradient values
85 state['exp_avg_sq'] = torch.zeros_like(p.data)
86
87 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
88 beta1, beta2 = group['betas']
89
90 state['step'] += 1
91
92 # Decay the first and second moment running average coefficient
93 # m_t
94 exp_avg.mul_(beta1).add_(1 - beta1, grad)
95 # v_t
96 exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
97
98 # Paper v3 does not use debiasing.
99 # bias_correction1 = 1 - beta1 ** state['step']
100 # bias_correction2 = 1 - beta2 ** state['step']
101 # Apply bias to lr to avoid broadcast.
102 step_size = group['lr'] # * math.sqrt(bias_correction2) / bias_correction1
103
104 weight_norm = p.data.pow(2).sum().sqrt().clamp(0, 10)
105
106 adam_step = exp_avg / exp_avg_sq.sqrt().add(group['eps'])
107 if group['weight_decay'] != 0:
108 adam_step.add_(group['weight_decay'], p.data)
109
110 adam_norm = adam_step.pow(2).sum().sqrt()
111 if weight_norm == 0 or adam_norm == 0:
112 trust_ratio = 1
113 else:
114 trust_ratio = weight_norm / adam_norm
115 state['weight_norm'] = weight_norm
116 state['adam_norm'] = adam_norm
117 state['trust_ratio'] = trust_ratio
118 if self.adam:
119 trust_ratio = 1
120
121 p.data.add_(-step_size * trust_ratio, adam_step)
122
123 return loss
Member Data Documentation
◆ adam
The documentation for this class was generated from the following file:
- smreid/fastreid/solver/optim/lamb.py
