release code
This commit is contained in:
miunangel
2
Dassl.ProGrad.pytorch/dassl/optim/__init__.py
Normal file
2
Dassl.ProGrad.pytorch/dassl/optim/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from .optimizer import build_optimizer
|
||||
from .lr_scheduler import build_lr_scheduler
|
||||
154
Dassl.ProGrad.pytorch/dassl/optim/lr_scheduler.py
Normal file
154
Dassl.ProGrad.pytorch/dassl/optim/lr_scheduler.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
Modified from https://github.com/KaiyangZhou/deep-person-reid
|
||||
"""
|
||||
import torch
|
||||
from torch.optim.lr_scheduler import _LRScheduler
|
||||
|
||||
AVAI_SCHEDS = ["single_step", "multi_step", "cosine"]
|
||||
|
||||
|
||||
class _BaseWarmupScheduler(_LRScheduler):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
successor,
|
||||
warmup_epoch,
|
||||
last_epoch=-1,
|
||||
verbose=False
|
||||
):
|
||||
self.successor = successor
|
||||
self.warmup_epoch = warmup_epoch
|
||||
super().__init__(optimizer, last_epoch, verbose)
|
||||
|
||||
def get_lr(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def step(self, epoch=None):
|
||||
if self.last_epoch >= self.warmup_epoch:
|
||||
self.successor.step(epoch)
|
||||
self._last_lr = self.successor.get_last_lr()
|
||||
else:
|
||||
super().step(epoch)
|
||||
|
||||
|
||||
class ConstantWarmupScheduler(_BaseWarmupScheduler):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
successor,
|
||||
warmup_epoch,
|
||||
cons_lr,
|
||||
last_epoch=-1,
|
||||
verbose=False
|
||||
):
|
||||
self.cons_lr = cons_lr
|
||||
super().__init__(
|
||||
optimizer, successor, warmup_epoch, last_epoch, verbose
|
||||
)
|
||||
|
||||
def get_lr(self):
|
||||
if self.last_epoch >= self.warmup_epoch:
|
||||
return self.successor.get_last_lr()
|
||||
return [self.cons_lr for _ in self.base_lrs]
|
||||
|
||||
|
||||
class LinearWarmupScheduler(_BaseWarmupScheduler):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
successor,
|
||||
warmup_epoch,
|
||||
min_lr,
|
||||
last_epoch=-1,
|
||||
verbose=False
|
||||
):
|
||||
self.min_lr = min_lr
|
||||
super().__init__(
|
||||
optimizer, successor, warmup_epoch, last_epoch, verbose
|
||||
)
|
||||
|
||||
def get_lr(self):
|
||||
if self.last_epoch >= self.warmup_epoch:
|
||||
return self.successor.get_last_lr()
|
||||
if self.last_epoch == 0:
|
||||
return [self.min_lr for _ in self.base_lrs]
|
||||
return [
|
||||
lr * self.last_epoch / self.warmup_epoch for lr in self.base_lrs
|
||||
]
|
||||
|
||||
|
||||
def build_lr_scheduler(optimizer, optim_cfg):
|
||||
"""A function wrapper for building a learning rate scheduler.
|
||||
|
||||
Args:
|
||||
optimizer (Optimizer): an Optimizer.
|
||||
optim_cfg (CfgNode): optimization config.
|
||||
"""
|
||||
lr_scheduler = optim_cfg.LR_SCHEDULER
|
||||
stepsize = optim_cfg.STEPSIZE
|
||||
gamma = optim_cfg.GAMMA
|
||||
max_epoch = optim_cfg.MAX_EPOCH
|
||||
|
||||
if lr_scheduler not in AVAI_SCHEDS:
|
||||
raise ValueError(
|
||||
"Unsupported scheduler: {}. Must be one of {}".format(
|
||||
lr_scheduler, AVAI_SCHEDS
|
||||
)
|
||||
)
|
||||
|
||||
if lr_scheduler == "single_step":
|
||||
if isinstance(stepsize, (list, tuple)):
|
||||
stepsize = stepsize[-1]
|
||||
|
||||
if not isinstance(stepsize, int):
|
||||
raise TypeError(
|
||||
"For single_step lr_scheduler, stepsize must "
|
||||
"be an integer, but got {}".format(type(stepsize))
|
||||
)
|
||||
|
||||
if stepsize <= 0:
|
||||
stepsize = max_epoch
|
||||
|
||||
scheduler = torch.optim.lr_scheduler.StepLR(
|
||||
optimizer, step_size=stepsize, gamma=gamma
|
||||
)
|
||||
|
||||
elif lr_scheduler == "multi_step":
|
||||
if not isinstance(stepsize, (list, tuple)):
|
||||
raise TypeError(
|
||||
"For multi_step lr_scheduler, stepsize must "
|
||||
"be a list, but got {}".format(type(stepsize))
|
||||
)
|
||||
|
||||
scheduler = torch.optim.lr_scheduler.MultiStepLR(
|
||||
optimizer, milestones=stepsize, gamma=gamma
|
||||
)
|
||||
|
||||
elif lr_scheduler == "cosine":
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||||
optimizer, float(max_epoch)
|
||||
)
|
||||
|
||||
if optim_cfg.WARMUP_EPOCH > 0:
|
||||
if not optim_cfg.WARMUP_RECOUNT:
|
||||
scheduler.last_epoch = optim_cfg.WARMUP_EPOCH
|
||||
|
||||
if optim_cfg.WARMUP_TYPE == "constant":
|
||||
scheduler = ConstantWarmupScheduler(
|
||||
optimizer, scheduler, optim_cfg.WARMUP_EPOCH,
|
||||
optim_cfg.WARMUP_CONS_LR
|
||||
)
|
||||
|
||||
elif optim_cfg.WARMUP_TYPE == "linear":
|
||||
scheduler = LinearWarmupScheduler(
|
||||
optimizer, scheduler, optim_cfg.WARMUP_EPOCH,
|
||||
optim_cfg.WARMUP_MIN_LR
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
return scheduler
|
||||
136
Dassl.ProGrad.pytorch/dassl/optim/optimizer.py
Normal file
136
Dassl.ProGrad.pytorch/dassl/optim/optimizer.py
Normal file
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
Modified from https://github.com/KaiyangZhou/deep-person-reid
|
||||
"""
|
||||
import warnings
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .radam import RAdam
|
||||
|
||||
AVAI_OPTIMS = ["adam", "amsgrad", "sgd", "rmsprop", "radam", "adamw"]
|
||||
|
||||
|
||||
def build_optimizer(model, optim_cfg):
|
||||
"""A function wrapper for building an optimizer.
|
||||
|
||||
Args:
|
||||
model (nn.Module or iterable): model.
|
||||
optim_cfg (CfgNode): optimization config.
|
||||
"""
|
||||
optim = optim_cfg.NAME
|
||||
lr = optim_cfg.LR
|
||||
weight_decay = optim_cfg.WEIGHT_DECAY
|
||||
momentum = optim_cfg.MOMENTUM
|
||||
sgd_dampening = optim_cfg.SGD_DAMPNING
|
||||
sgd_nesterov = optim_cfg.SGD_NESTEROV
|
||||
rmsprop_alpha = optim_cfg.RMSPROP_ALPHA
|
||||
adam_beta1 = optim_cfg.ADAM_BETA1
|
||||
adam_beta2 = optim_cfg.ADAM_BETA2
|
||||
staged_lr = optim_cfg.STAGED_LR
|
||||
new_layers = optim_cfg.NEW_LAYERS
|
||||
base_lr_mult = optim_cfg.BASE_LR_MULT
|
||||
|
||||
if optim not in AVAI_OPTIMS:
|
||||
raise ValueError(
|
||||
"Unsupported optim: {}. Must be one of {}".format(
|
||||
optim, AVAI_OPTIMS
|
||||
)
|
||||
)
|
||||
|
||||
if staged_lr:
|
||||
if not isinstance(model, nn.Module):
|
||||
raise TypeError(
|
||||
"When staged_lr is True, model given to "
|
||||
"build_optimizer() must be an instance of nn.Module"
|
||||
)
|
||||
|
||||
if isinstance(model, nn.DataParallel):
|
||||
model = model.module
|
||||
|
||||
if isinstance(new_layers, str):
|
||||
if new_layers is None:
|
||||
warnings.warn(
|
||||
"new_layers is empty, therefore, staged_lr is useless"
|
||||
)
|
||||
new_layers = [new_layers]
|
||||
|
||||
base_params = []
|
||||
base_layers = []
|
||||
new_params = []
|
||||
|
||||
for name, module in model.named_children():
|
||||
if name in new_layers:
|
||||
new_params += [p for p in module.parameters()]
|
||||
else:
|
||||
base_params += [p for p in module.parameters()]
|
||||
base_layers.append(name)
|
||||
|
||||
param_groups = [
|
||||
{
|
||||
"params": base_params,
|
||||
"lr": lr * base_lr_mult
|
||||
},
|
||||
{
|
||||
"params": new_params
|
||||
},
|
||||
]
|
||||
|
||||
else:
|
||||
if isinstance(model, nn.Module):
|
||||
param_groups = model.parameters()
|
||||
else:
|
||||
param_groups = model
|
||||
|
||||
if optim == "adam":
|
||||
optimizer = torch.optim.Adam(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
betas=(adam_beta1, adam_beta2),
|
||||
)
|
||||
|
||||
elif optim == "amsgrad":
|
||||
optimizer = torch.optim.Adam(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
betas=(adam_beta1, adam_beta2),
|
||||
amsgrad=True,
|
||||
)
|
||||
|
||||
elif optim == "sgd":
|
||||
optimizer = torch.optim.SGD(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
momentum=momentum,
|
||||
weight_decay=weight_decay,
|
||||
dampening=sgd_dampening,
|
||||
nesterov=sgd_nesterov,
|
||||
)
|
||||
|
||||
elif optim == "rmsprop":
|
||||
optimizer = torch.optim.RMSprop(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
momentum=momentum,
|
||||
weight_decay=weight_decay,
|
||||
alpha=rmsprop_alpha,
|
||||
)
|
||||
|
||||
elif optim == "radam":
|
||||
optimizer = RAdam(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
betas=(adam_beta1, adam_beta2),
|
||||
)
|
||||
|
||||
elif optim == "adamw":
|
||||
optimizer = torch.optim.AdamW(
|
||||
param_groups,
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
betas=(adam_beta1, adam_beta2),
|
||||
)
|
||||
|
||||
return optimizer
|
||||
332
Dassl.ProGrad.pytorch/dassl/optim/radam.py
Normal file
332
Dassl.ProGrad.pytorch/dassl/optim/radam.py
Normal file
@@ -0,0 +1,332 @@
|
||||
"""
|
||||
Imported from: https://github.com/LiyuanLucasLiu/RAdam
|
||||
|
||||
https://arxiv.org/abs/1908.03265
|
||||
|
||||
@article{liu2019radam,
|
||||
title={On the Variance of the Adaptive Learning Rate and Beyond},
|
||||
author={Liu, Liyuan and Jiang, Haoming and He, Pengcheng and Chen, Weizhu and Liu, Xiaodong and Gao, Jianfeng and Han, Jiawei},
|
||||
journal={arXiv preprint arXiv:1908.03265},
|
||||
year={2019}
|
||||
}
|
||||
"""
|
||||
import math
|
||||
import torch
|
||||
from torch.optim.optimizer import Optimizer
|
||||
|
||||
|
||||
class RAdam(Optimizer):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params,
|
||||
lr=1e-3,
|
||||
betas=(0.9, 0.999),
|
||||
eps=1e-8,
|
||||
weight_decay=0,
|
||||
degenerated_to_sgd=True,
|
||||
):
|
||||
if not 0.0 <= lr:
|
||||
raise ValueError("Invalid learning rate: {}".format(lr))
|
||||
if not 0.0 <= eps:
|
||||
raise ValueError("Invalid epsilon value: {}".format(eps))
|
||||
if not 0.0 <= betas[0] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 0: {}".format(betas[0])
|
||||
)
|
||||
if not 0.0 <= betas[1] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 1: {}".format(betas[1])
|
||||
)
|
||||
|
||||
self.degenerated_to_sgd = degenerated_to_sgd
|
||||
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
|
||||
self.buffer = [[None, None, None] for ind in range(10)]
|
||||
super(RAdam, self).__init__(params, defaults)
|
||||
|
||||
def __setstate__(self, state):
|
||||
super(RAdam, self).__setstate__(state)
|
||||
|
||||
def step(self, closure=None):
|
||||
|
||||
loss = None
|
||||
if closure is not None:
|
||||
loss = closure()
|
||||
|
||||
for group in self.param_groups:
|
||||
|
||||
for p in group["params"]:
|
||||
if p.grad is None:
|
||||
continue
|
||||
grad = p.grad.data.float()
|
||||
if grad.is_sparse:
|
||||
raise RuntimeError(
|
||||
"RAdam does not support sparse gradients"
|
||||
)
|
||||
|
||||
p_data_fp32 = p.data.float()
|
||||
|
||||
state = self.state[p]
|
||||
|
||||
if len(state) == 0:
|
||||
state["step"] = 0
|
||||
state["exp_avg"] = torch.zeros_like(p_data_fp32)
|
||||
state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
|
||||
else:
|
||||
state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
|
||||
state["exp_avg_sq"] = state["exp_avg_sq"].type_as(
|
||||
p_data_fp32
|
||||
)
|
||||
|
||||
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
||||
beta1, beta2 = group["betas"]
|
||||
|
||||
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
|
||||
exp_avg.mul_(beta1).add_(1 - beta1, grad)
|
||||
|
||||
state["step"] += 1
|
||||
buffered = self.buffer[int(state["step"] % 10)]
|
||||
if state["step"] == buffered[0]:
|
||||
N_sma, step_size = buffered[1], buffered[2]
|
||||
else:
|
||||
buffered[0] = state["step"]
|
||||
beta2_t = beta2**state["step"]
|
||||
N_sma_max = 2 / (1-beta2) - 1
|
||||
N_sma = N_sma_max - 2 * state["step"
|
||||
] * beta2_t / (1-beta2_t)
|
||||
buffered[1] = N_sma
|
||||
|
||||
# more conservative since it's an approximated value
|
||||
if N_sma >= 5:
|
||||
step_size = math.sqrt(
|
||||
(1-beta2_t) * (N_sma-4) / (N_sma_max-4) *
|
||||
(N_sma-2) / N_sma * N_sma_max / (N_sma_max-2)
|
||||
) / (1 - beta1**state["step"])
|
||||
elif self.degenerated_to_sgd:
|
||||
step_size = 1.0 / (1 - beta1**state["step"])
|
||||
else:
|
||||
step_size = -1
|
||||
buffered[2] = step_size
|
||||
|
||||
# more conservative since it's an approximated value
|
||||
if N_sma >= 5:
|
||||
if group["weight_decay"] != 0:
|
||||
p_data_fp32.add_(
|
||||
-group["weight_decay"] * group["lr"], p_data_fp32
|
||||
)
|
||||
denom = exp_avg_sq.sqrt().add_(group["eps"])
|
||||
p_data_fp32.addcdiv_(
|
||||
-step_size * group["lr"], exp_avg, denom
|
||||
)
|
||||
p.data.copy_(p_data_fp32)
|
||||
elif step_size > 0:
|
||||
if group["weight_decay"] != 0:
|
||||
p_data_fp32.add_(
|
||||
-group["weight_decay"] * group["lr"], p_data_fp32
|
||||
)
|
||||
p_data_fp32.add_(-step_size * group["lr"], exp_avg)
|
||||
p.data.copy_(p_data_fp32)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
class PlainRAdam(Optimizer):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params,
|
||||
lr=1e-3,
|
||||
betas=(0.9, 0.999),
|
||||
eps=1e-8,
|
||||
weight_decay=0,
|
||||
degenerated_to_sgd=True,
|
||||
):
|
||||
if not 0.0 <= lr:
|
||||
raise ValueError("Invalid learning rate: {}".format(lr))
|
||||
if not 0.0 <= eps:
|
||||
raise ValueError("Invalid epsilon value: {}".format(eps))
|
||||
if not 0.0 <= betas[0] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 0: {}".format(betas[0])
|
||||
)
|
||||
if not 0.0 <= betas[1] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 1: {}".format(betas[1])
|
||||
)
|
||||
|
||||
self.degenerated_to_sgd = degenerated_to_sgd
|
||||
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
|
||||
|
||||
super(PlainRAdam, self).__init__(params, defaults)
|
||||
|
||||
def __setstate__(self, state):
|
||||
super(PlainRAdam, self).__setstate__(state)
|
||||
|
||||
def step(self, closure=None):
|
||||
|
||||
loss = None
|
||||
if closure is not None:
|
||||
loss = closure()
|
||||
|
||||
for group in self.param_groups:
|
||||
|
||||
for p in group["params"]:
|
||||
if p.grad is None:
|
||||
continue
|
||||
grad = p.grad.data.float()
|
||||
if grad.is_sparse:
|
||||
raise RuntimeError(
|
||||
"RAdam does not support sparse gradients"
|
||||
)
|
||||
|
||||
p_data_fp32 = p.data.float()
|
||||
|
||||
state = self.state[p]
|
||||
|
||||
if len(state) == 0:
|
||||
state["step"] = 0
|
||||
state["exp_avg"] = torch.zeros_like(p_data_fp32)
|
||||
state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
|
||||
else:
|
||||
state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
|
||||
state["exp_avg_sq"] = state["exp_avg_sq"].type_as(
|
||||
p_data_fp32
|
||||
)
|
||||
|
||||
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
||||
beta1, beta2 = group["betas"]
|
||||
|
||||
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
|
||||
exp_avg.mul_(beta1).add_(1 - beta1, grad)
|
||||
|
||||
state["step"] += 1
|
||||
beta2_t = beta2**state["step"]
|
||||
N_sma_max = 2 / (1-beta2) - 1
|
||||
N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1-beta2_t)
|
||||
|
||||
# more conservative since it's an approximated value
|
||||
if N_sma >= 5:
|
||||
if group["weight_decay"] != 0:
|
||||
p_data_fp32.add_(
|
||||
-group["weight_decay"] * group["lr"], p_data_fp32
|
||||
)
|
||||
step_size = (
|
||||
group["lr"] * math.sqrt(
|
||||
(1-beta2_t) * (N_sma-4) / (N_sma_max-4) *
|
||||
(N_sma-2) / N_sma * N_sma_max / (N_sma_max-2)
|
||||
) / (1 - beta1**state["step"])
|
||||
)
|
||||
denom = exp_avg_sq.sqrt().add_(group["eps"])
|
||||
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
|
||||
p.data.copy_(p_data_fp32)
|
||||
elif self.degenerated_to_sgd:
|
||||
if group["weight_decay"] != 0:
|
||||
p_data_fp32.add_(
|
||||
-group["weight_decay"] * group["lr"], p_data_fp32
|
||||
)
|
||||
step_size = group["lr"] / (1 - beta1**state["step"])
|
||||
p_data_fp32.add_(-step_size, exp_avg)
|
||||
p.data.copy_(p_data_fp32)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
class AdamW(Optimizer):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params,
|
||||
lr=1e-3,
|
||||
betas=(0.9, 0.999),
|
||||
eps=1e-8,
|
||||
weight_decay=0,
|
||||
warmup=0
|
||||
):
|
||||
if not 0.0 <= lr:
|
||||
raise ValueError("Invalid learning rate: {}".format(lr))
|
||||
if not 0.0 <= eps:
|
||||
raise ValueError("Invalid epsilon value: {}".format(eps))
|
||||
if not 0.0 <= betas[0] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 0: {}".format(betas[0])
|
||||
)
|
||||
if not 0.0 <= betas[1] < 1.0:
|
||||
raise ValueError(
|
||||
"Invalid beta parameter at index 1: {}".format(betas[1])
|
||||
)
|
||||
|
||||
defaults = dict(
|
||||
lr=lr,
|
||||
betas=betas,
|
||||
eps=eps,
|
||||
weight_decay=weight_decay,
|
||||
warmup=warmup
|
||||
)
|
||||
super(AdamW, self).__init__(params, defaults)
|
||||
|
||||
def __setstate__(self, state):
|
||||
super(AdamW, self).__setstate__(state)
|
||||
|
||||
def step(self, closure=None):
|
||||
loss = None
|
||||
if closure is not None:
|
||||
loss = closure()
|
||||
|
||||
for group in self.param_groups:
|
||||
|
||||
for p in group["params"]:
|
||||
if p.grad is None:
|
||||
continue
|
||||
grad = p.grad.data.float()
|
||||
if grad.is_sparse:
|
||||
raise RuntimeError(
|
||||
"Adam does not support sparse gradients, please consider SparseAdam instead"
|
||||
)
|
||||
|
||||
p_data_fp32 = p.data.float()
|
||||
|
||||
state = self.state[p]
|
||||
|
||||
if len(state) == 0:
|
||||
state["step"] = 0
|
||||
state["exp_avg"] = torch.zeros_like(p_data_fp32)
|
||||
state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
|
||||
else:
|
||||
state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
|
||||
state["exp_avg_sq"] = state["exp_avg_sq"].type_as(
|
||||
p_data_fp32
|
||||
)
|
||||
|
||||
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
|
||||
beta1, beta2 = group["betas"]
|
||||
|
||||
state["step"] += 1
|
||||
|
||||
exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
|
||||
exp_avg.mul_(beta1).add_(1 - beta1, grad)
|
||||
|
||||
denom = exp_avg_sq.sqrt().add_(group["eps"])
|
||||
bias_correction1 = 1 - beta1**state["step"]
|
||||
bias_correction2 = 1 - beta2**state["step"]
|
||||
|
||||
if group["warmup"] > state["step"]:
|
||||
scheduled_lr = 1e-8 + state["step"] * group["lr"] / group[
|
||||
"warmup"]
|
||||
else:
|
||||
scheduled_lr = group["lr"]
|
||||
|
||||
step_size = (
|
||||
scheduled_lr * math.sqrt(bias_correction2) /
|
||||
bias_correction1
|
||||
)
|
||||
|
||||
if group["weight_decay"] != 0:
|
||||
p_data_fp32.add_(
|
||||
-group["weight_decay"] * scheduled_lr, p_data_fp32
|
||||
)
|
||||
|
||||
p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
|
||||
|
||||
p.data.copy_(p_data_fp32)
|
||||
|
||||
return loss
|
||||
Reference in New Issue
Block a user