DoruC/Grounded-Segment-Anything
0
1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch optimization for BERT model."""16 17import math18import warnings19from functools import partial20from typing import Callable, Iterable, Optional, Tuple, Union21 22import torch23from torch import nn24from torch.optim import Optimizer25from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau26 27from .trainer_utils import SchedulerType28from .utils import logging29from .utils.versions import require_version30 31 32logger = logging.get_logger(__name__)33 34 35def _get_constant_lambda(_=None):36 return 137 38 39def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):40 """41 Create a schedule with a constant learning rate, using the learning rate set in optimizer.42 43 Args:44 optimizer ([`~torch.optim.Optimizer`]):45 The optimizer for which to schedule the learning rate.46 last_epoch (`int`, *optional*, defaults to -1):47 The index of the last epoch when resuming training.48 49 Return:50 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.51 """52 53 return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch)54 55 56def get_reduce_on_plateau_schedule(optimizer: Optimizer):57 """58 Create a schedule with a constant learning rate that decreases when a metric has stopped improving.59 60 Args:61 optimizer ([`~torch.optim.Optimizer`]):62 The optimizer for which to schedule the learning rate.63 64 Return:65 `torch.optim.lr_scheduler.ReduceLROnPlateau` with the appropriate schedule.66 """67 68 return ReduceLROnPlateau(optimizer)69 70 71def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int):72 if current_step < num_warmup_steps:73 return float(current_step) / float(max(1.0, num_warmup_steps))74 return 1.075 76 77def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):78 """79 Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate80 increases linearly between 0 and the initial lr set in the optimizer.81 82 Args:83 optimizer ([`~torch.optim.Optimizer`]):84 The optimizer for which to schedule the learning rate.85 num_warmup_steps (`int`):86 The number of steps for the warmup phase.87 last_epoch (`int`, *optional*, defaults to -1):88 The index of the last epoch when resuming training.89 90 Return:91 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.92 """93 94 lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps)95 return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)96 97 98def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int):99 if current_step < num_warmup_steps:100 return float(current_step) / float(max(1, num_warmup_steps))101 return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))102 103 104def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):105 """106 Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after107 a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.108 109 Args:110 optimizer ([`~torch.optim.Optimizer`]):111 The optimizer for which to schedule the learning rate.112 num_warmup_steps (`int`):113 The number of steps for the warmup phase.114 num_training_steps (`int`):115 The total number of training steps.116 last_epoch (`int`, *optional*, defaults to -1):117 The index of the last epoch when resuming training.118 119 Return:120 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.121 """122 123 lr_lambda = partial(124 _get_linear_schedule_with_warmup_lr_lambda,125 num_warmup_steps=num_warmup_steps,126 num_training_steps=num_training_steps,127 )128 return LambdaLR(optimizer, lr_lambda, last_epoch)129 130 131def _get_cosine_schedule_with_warmup_lr_lambda(132 current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float133):134 if current_step < num_warmup_steps:135 return float(current_step) / float(max(1, num_warmup_steps))136 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))137 return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))138 139 140def get_cosine_schedule_with_warmup(141 optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1142):143 """144 Create a schedule with a learning rate that decreases following the values of the cosine function between the145 initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the146 initial lr set in the optimizer.147 148 Args:149 optimizer ([`~torch.optim.Optimizer`]):150 The optimizer for which to schedule the learning rate.151 num_warmup_steps (`int`):152 The number of steps for the warmup phase.153 num_training_steps (`int`):154 The total number of training steps.155 num_cycles (`float`, *optional*, defaults to 0.5):156 The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0157 following a half-cosine).158 last_epoch (`int`, *optional*, defaults to -1):159 The index of the last epoch when resuming training.160 161 Return:162 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.163 """164 165 lr_lambda = partial(166 _get_cosine_schedule_with_warmup_lr_lambda,167 num_warmup_steps=num_warmup_steps,168 num_training_steps=num_training_steps,169 num_cycles=num_cycles,170 )171 return LambdaLR(optimizer, lr_lambda, last_epoch)172 173 174def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(175 current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int176):177 if current_step < num_warmup_steps:178 return float(current_step) / float(max(1, num_warmup_steps))179 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))180 if progress >= 1.0:181 return 0.0182 return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))183 184 185def get_cosine_with_hard_restarts_schedule_with_warmup(186 optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1187):188 """189 Create a schedule with a learning rate that decreases following the values of the cosine function between the190 initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases191 linearly between 0 and the initial lr set in the optimizer.192 193 Args:194 optimizer ([`~torch.optim.Optimizer`]):195 The optimizer for which to schedule the learning rate.196 num_warmup_steps (`int`):197 The number of steps for the warmup phase.198 num_training_steps (`int`):199 The total number of training steps.200 num_cycles (`int`, *optional*, defaults to 1):201 The number of hard restarts to use.202 last_epoch (`int`, *optional*, defaults to -1):203 The index of the last epoch when resuming training.204 205 Return:206 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.207 """208 209 lr_lambda = partial(210 _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda,211 num_warmup_steps=num_warmup_steps,212 num_training_steps=num_training_steps,213 num_cycles=num_cycles,214 )215 return LambdaLR(optimizer, lr_lambda, last_epoch)216 217 218def _get_polynomial_decay_schedule_with_warmup_lr_lambda(219 current_step: int,220 *,221 num_warmup_steps: int,222 num_training_steps: int,223 lr_end: float,224 power: float,225 lr_init: int,226):227 if current_step < num_warmup_steps:228 return float(current_step) / float(max(1, num_warmup_steps))229 elif current_step > num_training_steps:230 return lr_end / lr_init # as LambdaLR multiplies by lr_init231 else:232 lr_range = lr_init - lr_end233 decay_steps = num_training_steps - num_warmup_steps234 pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps235 decay = lr_range * pct_remaining**power + lr_end236 return decay / lr_init # as LambdaLR multiplies by lr_init237 238 239def get_polynomial_decay_schedule_with_warmup(240 optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1241):242 """243 Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the244 optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the245 initial lr set in the optimizer.246 247 Args:248 optimizer ([`~torch.optim.Optimizer`]):249 The optimizer for which to schedule the learning rate.250 num_warmup_steps (`int`):251 The number of steps for the warmup phase.252 num_training_steps (`int`):253 The total number of training steps.254 lr_end (`float`, *optional*, defaults to 1e-7):255 The end LR.256 power (`float`, *optional*, defaults to 1.0):257 Power factor.258 last_epoch (`int`, *optional*, defaults to -1):259 The index of the last epoch when resuming training.260 261 Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT262 implementation at263 https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37264 265 Return:266 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.267 268 """269 270 lr_init = optimizer.defaults["lr"]271 if not (lr_init > lr_end):272 raise ValueError(f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})")273 274 lr_lambda = partial(275 _get_polynomial_decay_schedule_with_warmup_lr_lambda,276 num_warmup_steps=num_warmup_steps,277 num_training_steps=num_training_steps,278 lr_end=lr_end,279 power=power,280 lr_init=lr_init,281 )282 return LambdaLR(optimizer, lr_lambda, last_epoch)283 284 285def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: int = None):286 if current_step < num_warmup_steps:287 return float(current_step) / float(max(1, num_warmup_steps))288 shift = timescale - num_warmup_steps289 decay = 1.0 / math.sqrt((current_step + shift) / timescale)290 return decay291 292 293def get_inverse_sqrt_schedule(294 optimizer: Optimizer, num_warmup_steps: int, timescale: int = None, last_epoch: int = -1295):296 """297 Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a298 warmup period which increases lr linearly from 0 to the initial lr set in the optimizer.299 300 Args:301 optimizer ([`~torch.optim.Optimizer`]):302 The optimizer for which to schedule the learning rate.303 num_warmup_steps (`int`):304 The number of steps for the warmup phase.305 timescale (`int`, *optional*, defaults to `num_warmup_steps`):306 Time scale.307 last_epoch (`int`, *optional*, defaults to -1):308 The index of the last epoch when resuming training.309 310 Return:311 `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.312 """313 # Note: this implementation is adapted from314 # https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930315 316 if timescale is None:317 timescale = num_warmup_steps318 319 lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale)320 return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)321 322 323TYPE_TO_SCHEDULER_FUNCTION = {324 SchedulerType.LINEAR: get_linear_schedule_with_warmup,325 SchedulerType.COSINE: get_cosine_schedule_with_warmup,326 SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,327 SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,328 SchedulerType.CONSTANT: get_constant_schedule,329 SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,330 SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule,331 SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule,332}333 334 335def get_scheduler(336 name: Union[str, SchedulerType],337 optimizer: Optimizer,338 num_warmup_steps: Optional[int] = None,339 num_training_steps: Optional[int] = None,340):341 """342 Unified API to get any scheduler from its name.343 344 Args:345 name (`str` or `SchedulerType`):346 The name of the scheduler to use.347 optimizer (`torch.optim.Optimizer`):348 The optimizer that will be used during training.349 num_warmup_steps (`int`, *optional*):350 The number of warmup steps to do. This is not required by all schedulers (hence the argument being351 optional), the function will raise an error if it's unset and the scheduler type requires it.352 num_training_steps (`int``, *optional*):353 The number of training steps to do. This is not required by all schedulers (hence the argument being354 optional), the function will raise an error if it's unset and the scheduler type requires it.355 """356 name = SchedulerType(name)357 schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]358 if name == SchedulerType.CONSTANT or name == SchedulerType.REDUCE_ON_PLATEAU:359 return schedule_func(optimizer)360 361 # All other schedulers require `num_warmup_steps`362 if num_warmup_steps is None:363 raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")364 365 if name == SchedulerType.CONSTANT_WITH_WARMUP:366 return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)367 368 if name == SchedulerType.INVERSE_SQRT:369 return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)370 371 # All other schedulers require `num_training_steps`372 if num_training_steps is None:373 raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")374 375 return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)376 377 378class AdamW(Optimizer):379 """380 Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay381 Regularization](https://arxiv.org/abs/1711.05101).382 383 Parameters:384 params (`Iterable[nn.parameter.Parameter]`):385 Iterable of parameters to optimize or dictionaries defining parameter groups.386 lr (`float`, *optional*, defaults to 0.001):387 The learning rate to use.388 betas (`Tuple[float,float]`, *optional*, defaults to `(0.9, 0.999)`):389 Adam's betas parameters (b1, b2).390 eps (`float`, *optional*, defaults to 1e-06):391 Adam's epsilon for numerical stability.392 weight_decay (`float`, *optional*, defaults to 0.0):393 Decoupled weight decay to apply.394 correct_bias (`bool`, *optional*, defaults to `True`):395 Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).396 no_deprecation_warning (`bool`, *optional*, defaults to `False`):397 A flag used to disable the deprecation warning (set to `True` to disable the warning).398 """399 400 def __init__(401 self,402 params: Iterable[nn.parameter.Parameter],403 lr: float = 1e-3,404 betas: Tuple[float, float] = (0.9, 0.999),405 eps: float = 1e-6,406 weight_decay: float = 0.0,407 correct_bias: bool = True,408 no_deprecation_warning: bool = False,409 ):410 if not no_deprecation_warning:411 warnings.warn(412 "This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch"413 " implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this"414 " warning",415 FutureWarning,416 )417 require_version("torch>=1.5.0") # add_ with alpha418 if lr < 0.0:419 raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")420 if not 0.0 <= betas[0] < 1.0:421 raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")422 if not 0.0 <= betas[1] < 1.0:423 raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")424 if not 0.0 <= eps:425 raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")426 defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "correct_bias": correct_bias}427 super().__init__(params, defaults)428 429 @torch.no_grad()430 def step(self, closure: Callable = None):431 """432 Performs a single optimization step.433 434 Arguments:435 closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.436 """437 loss = None438 if closure is not None:439 loss = closure()440 441 for group in self.param_groups:442 for p in group["params"]:443 if p.grad is None:444 continue445 grad = p.grad446 if grad.is_sparse:447 raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead")448 449 state = self.state[p]450 451 # State initialization452 if len(state) == 0:453 state["step"] = 0454 # Exponential moving average of gradient values455 state["exp_avg"] = torch.zeros_like(p)456 # Exponential moving average of squared gradient values457 state["exp_avg_sq"] = torch.zeros_like(p)458 459 exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]460 beta1, beta2 = group["betas"]461 462 state["step"] += 1463 464 # Decay the first and second moment running average coefficient465 # In-place operations to update the averages at the same time466 exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))467 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)468 denom = exp_avg_sq.sqrt().add_(group["eps"])469 470 step_size = group["lr"]471 if group["correct_bias"]: # No bias correction for Bert472 bias_correction1 = 1.0 - beta1 ** state["step"]473 bias_correction2 = 1.0 - beta2 ** state["step"]474 step_size = step_size * math.sqrt(bias_correction2) / bias_correction1475 476 p.addcdiv_(exp_avg, denom, value=-step_size)477 478 # Just adding the square of the weights to the loss function is *not*479 # the correct way of using L2 regularization/weight decay with Adam,480 # since that will interact with the m and v parameters in strange ways.481 #482 # Instead we want to decay the weights in a manner that doesn't interact483 # with the m/v parameters. This is equivalent to adding the square484 # of the weights to the loss with plain (non-momentum) SGD.485 # Add weight decay at the end (fixed version)486 if group["weight_decay"] > 0.0:487 p.add_(p, alpha=(-group["lr"] * group["weight_decay"]))488 489 return loss490 491 492class Adafactor(Optimizer):493 """494 AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:495 https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py496 497 Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that498 this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and499 `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and500 `relative_step=False`.501 502 Arguments:503 params (`Iterable[nn.parameter.Parameter]`):504 Iterable of parameters to optimize or dictionaries defining parameter groups.505 lr (`float`, *optional*):506 The external learning rate.507 eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):508 Regularization constants for square gradient and parameter scale respectively509 clip_threshold (`float`, *optional*, defaults to 1.0):510 Threshold of root mean square of final gradient update511 decay_rate (`float`, *optional*, defaults to -0.8):512 Coefficient used to compute running averages of square513 beta1 (`float`, *optional*):514 Coefficient used for computing running averages of gradient515 weight_decay (`float`, *optional*, defaults to 0.0):516 Weight decay (L2 penalty)517 scale_parameter (`bool`, *optional*, defaults to `True`):518 If True, learning rate is scaled by root mean square519 relative_step (`bool`, *optional*, defaults to `True`):520 If True, time-dependent learning rate is computed instead of external learning rate521 warmup_init (`bool`, *optional*, defaults to `False`):522 Time-dependent learning rate computation depends on whether warm-up initialization is being used523 524 This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.525 526 Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):527 528 - Training without LR warmup or clip_threshold is not recommended.529 530 - use scheduled LR warm-up to fixed LR531 - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235)532 - Disable relative updates533 - Use scale_parameter=False534 - Additional optimizer operations like gradient clipping should not be used alongside Adafactor535 536 Example:537 538 ```python539 Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)540 ```541 542 Others reported the following combination to work well:543 544 ```python545 Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)546 ```547 548 When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]549 scheduler as following:550 551 ```python552 from transformers.optimization import Adafactor, AdafactorSchedule553 554 optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)555 lr_scheduler = AdafactorSchedule(optimizer)556 trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))557 ```558 559 Usage:560 561 ```python562 # replace AdamW with Adafactor563 optimizer = Adafactor(564 model.parameters(),565 lr=1e-3,566 eps=(1e-30, 1e-3),567 clip_threshold=1.0,568 decay_rate=-0.8,569 beta1=None,570 weight_decay=0.0,571 relative_step=False,572 scale_parameter=False,573 warmup_init=False,574 )575 ```"""576 577 def __init__(578 self,579 params,580 lr=None,581 eps=(1e-30, 1e-3),582 clip_threshold=1.0,583 decay_rate=-0.8,584 beta1=None,585 weight_decay=0.0,586 scale_parameter=True,587 relative_step=True,588 warmup_init=False,589 ):590 require_version("torch>=1.5.0") # add_ with alpha591 if lr is not None and relative_step:592 raise ValueError("Cannot combine manual `lr` and `relative_step=True` options")593 if warmup_init and not relative_step:594 raise ValueError("`warmup_init=True` requires `relative_step=True`")595 596 defaults = {597 "lr": lr,598 "eps": eps,599 "clip_threshold": clip_threshold,600 "decay_rate": decay_rate,601 "beta1": beta1,602 "weight_decay": weight_decay,603 "scale_parameter": scale_parameter,604 "relative_step": relative_step,605 "warmup_init": warmup_init,606 }607 super().__init__(params, defaults)608 609 @staticmethod610 def _get_lr(param_group, param_state):611 rel_step_sz = param_group["lr"]612 if param_group["relative_step"]:613 min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2614 rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))615 param_scale = 1.0616 if param_group["scale_parameter"]:617 param_scale = max(param_group["eps"][1], param_state["RMS"])618 return param_scale * rel_step_sz619 620 @staticmethod621 def _get_options(param_group, param_shape):622 factored = len(param_shape) >= 2623 use_first_moment = param_group["beta1"] is not None624 return factored, use_first_moment625 626 @staticmethod627 def _rms(tensor):628 return tensor.norm(2) / (tensor.numel() ** 0.5)629 630 @staticmethod631 def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):632 # copy from fairseq's adafactor implementation:633 # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505634 r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)635 c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()636 return torch.mul(r_factor, c_factor)637 638 @torch.no_grad()639 def step(self, closure=None):640 """641 Performs a single optimization step642 643 Arguments:644 closure (callable, optional): A closure that reevaluates the model645 and returns the loss.646 """647 loss = None648 if closure is not None:649 loss = closure()650 651 for group in self.param_groups:652 for p in group["params"]:653 if p.grad is None:654 continue655 grad = p.grad656 if grad.dtype in {torch.float16, torch.bfloat16}:657 grad = grad.float()658 if grad.is_sparse:659 raise RuntimeError("Adafactor does not support sparse gradients.")660 661 state = self.state[p]662 grad_shape = grad.shape663 664 factored, use_first_moment = self._get_options(group, grad_shape)665 # State Initialization666 if len(state) == 0:667 state["step"] = 0668 669 if use_first_moment:670 # Exponential moving average of gradient values671 state["exp_avg"] = torch.zeros_like(grad)672 if factored:673 state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)674 state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)675 else:676 state["exp_avg_sq"] = torch.zeros_like(grad)677 678 state["RMS"] = 0679 else:680 if use_first_moment:681 state["exp_avg"] = state["exp_avg"].to(grad)682 if factored:683 state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)684 state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)685 else:686 state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)687 688 p_data_fp32 = p689 if p.dtype in {torch.float16, torch.bfloat16}:690 p_data_fp32 = p_data_fp32.float()691 692 state["step"] += 1693 state["RMS"] = self._rms(p_data_fp32)694 lr = self._get_lr(group, state)695 696 beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])697 update = (grad**2) + group["eps"][0]698 if factored:699 exp_avg_sq_row = state["exp_avg_sq_row"]700 exp_avg_sq_col = state["exp_avg_sq_col"]701 702 exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))703 exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))704 705 # Approximation of exponential moving average of square of gradient706 update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)707 update.mul_(grad)708 else:709 exp_avg_sq = state["exp_avg_sq"]710 711 exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))712 update = exp_avg_sq.rsqrt().mul_(grad)713 714 update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))715 update.mul_(lr)716 717 if use_first_moment:718 exp_avg = state["exp_avg"]719 exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))720 update = exp_avg721 722 if group["weight_decay"] != 0:723 p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))724 725 p_data_fp32.add_(-update)726 727 if p.dtype in {torch.float16, torch.bfloat16}:728 p.copy_(p_data_fp32)729 730 return loss731 732 733class AdafactorSchedule(LambdaLR):734 """735 Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g.,736 for logging), this class creates a proxy object that retrieves the current lr values from the optimizer.737 738 It returns `initial_lr` during startup and the actual `lr` during stepping.739 """740 741 def __init__(self, optimizer, initial_lr=0.0):742 def lr_lambda(_):743 return initial_lr744 745 for group in optimizer.param_groups:746 group["initial_lr"] = initial_lr747 super().__init__(optimizer, lr_lambda)748 for group in optimizer.param_groups:749 del group["initial_lr"]750 751 def get_lr(self):752 opt = self.optimizer753 lrs = [754 opt._get_lr(group, opt.state[group["params"][0]])755 for group in opt.param_groups756 if group["params"][0].grad is not None757 ]758 if len(lrs) == 0:759 lrs = self.base_lrs # if called before stepping760 return lrs761 762 763def get_adafactor_schedule(optimizer, initial_lr=0.0):764 """765 Get a proxy schedule for [`~optimization.Adafactor`]766 767 Args:768 optimizer ([`~torch.optim.Optimizer`]):769 The optimizer for which to schedule the learning rate.770 initial_lr (`float`, *optional*, defaults to 0.0):771 Initial lr772 773 Return:774 [`~optimization.Adafactor`] proxy schedule object.775 776 777 """778 return AdafactorSchedule(optimizer, initial_lr)779