File size: 3,110 Bytes
1197f7d 3e08dd8 1197f7d 3e08dd8 1197f7d dcceddd 1197f7d dcceddd 1197f7d 3e08dd8 1197f7d dcceddd 1197f7d 3e08dd8 97e9dcb 3e08dd8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from typing import Any, Dict, Type
import torch
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
from yolo.config.config import OptimizerConfig, SchedulerConfig
from yolo.model.yolo import YOLO
class ExponentialMovingAverage:
def __init__(self, model: torch.nn.Module, decay: float):
self.model = model
self.decay = decay
self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}
def update(self):
"""Update the shadow parameters using the current model parameters."""
for name, param in self.model.named_parameters():
assert name in self.shadow, "All model parameters should have a corresponding shadow parameter."
new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
self.shadow[name] = new_average.clone()
def apply_shadow(self):
"""Apply the shadow parameters to the model."""
for name, param in self.model.named_parameters():
param.data.copy_(self.shadow[name])
def restore(self):
"""Restore the original parameters from the shadow."""
for name, param in self.model.named_parameters():
self.shadow[name].copy_(param.data)
def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
"""Create an optimizer for the given model parameters based on the configuration.
Returns:
An instance of the optimizer configured according to the provided settings.
"""
optimizer_class: Type[Optimizer] = getattr(torch.optim, optim_cfg.type)
bias_params = [p for name, p in model.named_parameters() if "bias" in name]
norm_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" in name]
conv_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" not in name]
model_parameters = [
{"params": bias_params, "nestrov": True, "momentum": 0.937},
{"params": conv_params, "weight_decay": 0.0},
{"params": norm_params, "weight_decay": 1e-5},
]
return optimizer_class(model_parameters, **optim_cfg.args)
def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
"""Create a learning rate scheduler for the given optimizer based on the configuration.
Returns:
An instance of the scheduler configured according to the provided settings.
"""
scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedule_cfg.type)
schedule = scheduler_class(optimizer, **schedule_cfg.args)
if hasattr(schedule_cfg, "warmup"):
wepoch = schedule_cfg.warmup.epochs
lambda1 = lambda epoch: 0.1 + 0.9 * (epoch + 1 / wepoch) if epoch < wepoch else 1
lambda2 = lambda epoch: 10 - 9 * (epoch / wepoch) if epoch < wepoch else 1
warmup_schedule = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2, lambda1])
schedule = SequentialLR(optimizer, schedulers=[warmup_schedule, schedule], milestones=[2])
return schedule
|