Source code for holocron.optim.tadam

# Copyright (C) 2019-2022, François-Guillaume Fernandez.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.

from typing import Callable, Dict, Iterable, Optional, Tuple

import torch
from torch.optim.optimizer import Optimizer

from . import functional as F


[docs] class TAdam(Optimizer): """Implements the TAdam optimizer from `"TAdam: A Robust Stochastic Gradient Optimizer" <https://arxiv.org/pdf/2003.00179.pdf>`_. Args: params (iterable): iterable of parameters to optimize or dicts defining parameter groups lr (float, optional): learning rate betas (Tuple[float, float], optional): coefficients used for running averages (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) dof (int, optional): degrees of freedom """ def __init__( self, params: Iterable[torch.nn.Parameter], # type: ignore[name-defined] lr: float = 1e-3, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.0, amsgrad: bool = False, dof: Optional[float] = None, ) -> None: if lr < 0.0: raise ValueError(f"Invalid learning rate: {lr}") if eps < 0.0: raise ValueError(f"Invalid epsilon value: {eps}") if not 0.0 <= betas[0] < 1.0: raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") if not 0.0 <= betas[1] < 1.0: raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") if not 0.0 <= weight_decay: raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad, dof=dof) super().__init__(params, defaults) def __setstate__(self, state: Dict[str, torch.Tensor]) -> None: super().__setstate__(state) for group in self.param_groups: group.setdefault("amsgrad", False) @torch.no_grad() def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]: """Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: params_with_grad = [] grads = [] exp_avgs = [] exp_avg_sqs = [] W_ts = [] max_exp_avg_sqs = [] state_steps = [] beta1, beta2 = group["betas"] for p in group["params"]: if p.grad is not None: params_with_grad.append(p) if p.grad.is_sparse: raise RuntimeError(f"{self.__class__.__name__} does not support sparse gradients") grads.append(p.grad) state = self.state[p] # Lazy state initialization if len(state) == 0: state["step"] = 0 # Exponential moving average of gradient values state["exp_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format) # Exponential moving average of squared gradient values state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.preserve_format) if group["amsgrad"]: # Maintains max of all exp. moving avg. of sq. grad. values state["max_exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.preserve_format) # Tadam specific state["W_t"] = beta1 / (1 - beta1) * torch.ones(1, dtype=p.data.dtype, device=p.data.device) exp_avgs.append(state["exp_avg"]) exp_avg_sqs.append(state["exp_avg_sq"]) W_ts.append(state["W_t"]) if group["amsgrad"]: max_exp_avg_sqs.append(state["max_exp_avg_sq"]) # update the steps for each param group update state["step"] += 1 # record the step after step update state_steps.append(state["step"]) F.tadam( params_with_grad, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, W_ts, state_steps, group["amsgrad"], beta1, beta2, group["lr"], group["weight_decay"], group["eps"], group["dof"], ) return loss