Spaces:
Running
on
T4
Running
on
T4
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
import torch | |
import itertools | |
from .lr_scheduler import WarmupMultiStepLR, WarmupCosineAnnealingLR, WarmupReduceLROnPlateau | |
def make_optimizer(cfg, model): | |
def maybe_add_full_model_gradient_clipping(optim): # optim: the optimizer class | |
# detectron2 doesn't have full model gradient clipping now | |
clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE | |
enable = ( | |
cfg.SOLVER.CLIP_GRADIENTS.ENABLED | |
and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model" | |
and clip_norm_val > 0.0 | |
) | |
class FullModelGradientClippingOptimizer(optim): | |
def step(self, closure=None): | |
all_params = itertools.chain(*[x["params"] for x in self.param_groups]) | |
torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val) | |
super().step(closure=closure) | |
return FullModelGradientClippingOptimizer if enable else optim | |
params = [] | |
for key, value in model.named_parameters(): | |
if not value.requires_grad: | |
print(key, "no grad") | |
continue | |
lr = cfg.SOLVER.BASE_LR | |
weight_decay = cfg.SOLVER.WEIGHT_DECAY | |
# different lr schedule | |
if "language_backbone" in key: | |
lr = cfg.SOLVER.LANG_LR | |
if "backbone.body" in key and "language_backbone.body" not in key: | |
lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BACKBONE_BODY_LR_FACTOR | |
if "t2i" in key: | |
lr = lr * cfg.SOLVER.FUSION_LR_FACTOR | |
if "bias" in key: | |
lr *= cfg.SOLVER.BIAS_LR_FACTOR | |
weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS | |
if "norm" in key or "Norm" in key: | |
weight_decay *= cfg.SOLVER.WEIGHT_DECAY_NORM_FACTOR | |
print("Setting weight decay of {} to {}".format(key, weight_decay)) | |
print(key, lr, ) | |
params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}] | |
if cfg.SOLVER.OPTIMIZER == "SGD": | |
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(params, lr, momentum=cfg.SOLVER.MOMENTUM) | |
elif cfg.SOLVER.OPTIMIZER == "ADAMW": | |
optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(params, lr) | |
return optimizer | |
def make_lr_scheduler(cfg, optimizer): | |
if cfg.SOLVER.MULTI_MAX_EPOCH: | |
assert len(cfg.SOLVER.MULTI_MAX_EPOCH) == len(cfg.SOLVER.STEPS) | |
lr_scheduler = [] | |
for stage_step, stage_max_epoch in zip(cfg.SOLVER.STEPS, cfg.SOLVER.MULTI_MAX_ITER): | |
milestones = [] | |
for step in stage_step: | |
milestones.append(round(step * stage_max_epoch)) | |
lr_scheduler.append( | |
WarmupMultiStepLR( | |
optimizer, | |
milestones, | |
cfg.SOLVER.GAMMA, | |
warmup_factor=cfg.SOLVER.WARMUP_FACTOR, | |
warmup_iters=cfg.SOLVER.WARMUP_ITERS, | |
warmup_method=cfg.SOLVER.WARMUP_METHOD, | |
) | |
) | |
return lr_scheduler | |
elif cfg.SOLVER.USE_COSINE: | |
max_iters = cfg.SOLVER.MAX_ITER | |
return WarmupCosineAnnealingLR( | |
optimizer, | |
max_iters, | |
cfg.SOLVER.GAMMA, | |
warmup_factor=cfg.SOLVER.WARMUP_FACTOR, | |
warmup_iters=cfg.SOLVER.WARMUP_ITERS, | |
warmup_method=cfg.SOLVER.WARMUP_METHOD, | |
eta_min=cfg.SOLVER.MIN_LR, | |
) | |
elif cfg.SOLVER.USE_AUTOSTEP: | |
max_iters = cfg.SOLVER.MAX_ITER | |
return WarmupReduceLROnPlateau( | |
optimizer, | |
max_iters, | |
cfg.SOLVER.GAMMA, | |
warmup_factor=cfg.SOLVER.WARMUP_FACTOR, | |
warmup_iters=cfg.SOLVER.WARMUP_ITERS, | |
warmup_method=cfg.SOLVER.WARMUP_METHOD, | |
eta_min=cfg.SOLVER.MIN_LR, | |
patience=cfg.SOLVER.STEP_PATIENCE, | |
verbose=True, | |
) | |
else: | |
milestones = [] | |
for step in cfg.SOLVER.STEPS: | |
if step < 1: | |
milestones.append(round(step * cfg.SOLVER.MAX_ITER)) | |
else: | |
milestones.append(step) | |
return WarmupMultiStepLR( | |
optimizer, | |
milestones, | |
cfg.SOLVER.GAMMA, | |
warmup_factor=cfg.SOLVER.WARMUP_FACTOR, | |
warmup_iters=cfg.SOLVER.WARMUP_ITERS, | |
warmup_method=cfg.SOLVER.WARMUP_METHOD, | |
) | |