Implements Adam algorithm with weight decay fix. Parameters: lr (float): learning rate. Default 1e-3. betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999) eps (float): Adams epsilon. Default: 1e-6 weight_decay (float): Weight decay. D
| 94 | |
| 95 | |
| 96 | class AdamW(Optimizer): |
| 97 | """ Implements Adam algorithm with weight decay fix. |
| 98 | |
| 99 | Parameters: |
| 100 | lr (float): learning rate. Default 1e-3. |
| 101 | betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999) |
| 102 | eps (float): Adams epsilon. Default: 1e-6 |
| 103 | weight_decay (float): Weight decay. Default: 0.0 |
| 104 | correct_bias (bool): can be set to False to avoid correcting bias in Adam (e.g. like in Bert TF repository). Default True. |
| 105 | """ |
| 106 | |
| 107 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, correct_bias=True): |
| 108 | if lr < 0.0: |
| 109 | raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr)) |
| 110 | if not 0.0 <= betas[0] < 1.0: |
| 111 | raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[0])) |
| 112 | if not 0.0 <= betas[1] < 1.0: |
| 113 | raise ValueError("Invalid beta parameter: {} - should be in [0.0, 1.0[".format(betas[1])) |
| 114 | if not 0.0 <= eps: |
| 115 | raise ValueError("Invalid epsilon value: {} - should be >= 0.0".format(eps)) |
| 116 | defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias) |
| 117 | super().__init__(params, defaults) |
| 118 | |
| 119 | def step(self, closure=None): |
| 120 | """Performs a single optimization step. |
| 121 | |
| 122 | Arguments: |
| 123 | closure (callable, optional): A closure that reevaluates the model |
| 124 | and returns the loss. |
| 125 | """ |
| 126 | loss = None |
| 127 | if closure is not None: |
| 128 | loss = closure() |
| 129 | |
| 130 | for group in self.param_groups: |
| 131 | for p in group["params"]: |
| 132 | if p.grad is None: |
| 133 | continue |
| 134 | grad = p.grad.data |
| 135 | if grad.is_sparse: |
| 136 | raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead") |
| 137 | |
| 138 | state = self.state[p] |
| 139 | |
| 140 | # State initialization |
| 141 | if len(state) == 0: |
| 142 | state["step"] = 0 |
| 143 | # Exponential moving average of gradient values |
| 144 | state["exp_avg"] = torch.zeros_like(p.data) |
| 145 | # Exponential moving average of squared gradient values |
| 146 | state["exp_avg_sq"] = torch.zeros_like(p.data) |
| 147 | |
| 148 | exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] |
| 149 | beta1, beta2 = group["betas"] |
| 150 | |
| 151 | state["step"] += 1 |
| 152 | |
| 153 | # Decay the first and second moment running average coefficient |
no outgoing calls