Spaces:
Runtime error
Runtime error
File size: 4,611 Bytes
3094730 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
# Copyright (c) OpenMMLab. All rights reserved.
import math
from typing import Optional
import numpy as np
from mmengine.hooks import ParamSchedulerHook
from mmengine.runner import Runner
from mmyolo.registry import HOOKS
def linear_fn(lr_factor: float, max_epochs: int):
"""Generate linear function."""
return lambda x: (1 - x / max_epochs) * (1.0 - lr_factor) + lr_factor
def cosine_fn(lr_factor: float, max_epochs: int):
"""Generate cosine function."""
return lambda x: (
(1 - math.cos(x * math.pi / max_epochs)) / 2) * (lr_factor - 1) + 1
@HOOKS.register_module()
class YOLOv5ParamSchedulerHook(ParamSchedulerHook):
"""A hook to update learning rate and momentum in optimizer of YOLOv5."""
priority = 9
scheduler_maps = {'linear': linear_fn, 'cosine': cosine_fn}
def __init__(self,
scheduler_type: str = 'linear',
lr_factor: float = 0.01,
max_epochs: int = 300,
warmup_epochs: int = 3,
warmup_bias_lr: float = 0.1,
warmup_momentum: float = 0.8,
warmup_mim_iter: int = 1000,
**kwargs):
assert scheduler_type in self.scheduler_maps
self.warmup_epochs = warmup_epochs
self.warmup_bias_lr = warmup_bias_lr
self.warmup_momentum = warmup_momentum
self.warmup_mim_iter = warmup_mim_iter
kwargs.update({'lr_factor': lr_factor, 'max_epochs': max_epochs})
self.scheduler_fn = self.scheduler_maps[scheduler_type](**kwargs)
self._warmup_end = False
self._base_lr = None
self._base_momentum = None
def before_train(self, runner: Runner):
"""Operations before train.
Args:
runner (Runner): The runner of the training process.
"""
optimizer = runner.optim_wrapper.optimizer
for group in optimizer.param_groups:
# If the param is never be scheduled, record the current value
# as the initial value.
group.setdefault('initial_lr', group['lr'])
group.setdefault('initial_momentum', group.get('momentum', -1))
self._base_lr = [
group['initial_lr'] for group in optimizer.param_groups
]
self._base_momentum = [
group['initial_momentum'] for group in optimizer.param_groups
]
def before_train_iter(self,
runner: Runner,
batch_idx: int,
data_batch: Optional[dict] = None):
"""Operations before each training iteration.
Args:
runner (Runner): The runner of the training process.
batch_idx (int): The index of the current batch in the train loop.
data_batch (dict or tuple or list, optional): Data from dataloader.
"""
cur_iters = runner.iter
cur_epoch = runner.epoch
optimizer = runner.optim_wrapper.optimizer
# The minimum warmup is self.warmup_mim_iter
warmup_total_iters = max(
round(self.warmup_epochs * len(runner.train_dataloader)),
self.warmup_mim_iter)
if cur_iters <= warmup_total_iters:
xp = [0, warmup_total_iters]
for group_idx, param in enumerate(optimizer.param_groups):
if group_idx == 2:
# bias learning rate will be handled specially
yp = [
self.warmup_bias_lr,
self._base_lr[group_idx] * self.scheduler_fn(cur_epoch)
]
else:
yp = [
0.0,
self._base_lr[group_idx] * self.scheduler_fn(cur_epoch)
]
param['lr'] = np.interp(cur_iters, xp, yp)
if 'momentum' in param:
param['momentum'] = np.interp(
cur_iters, xp,
[self.warmup_momentum, self._base_momentum[group_idx]])
else:
self._warmup_end = True
def after_train_epoch(self, runner: Runner):
"""Operations after each training epoch.
Args:
runner (Runner): The runner of the training process.
"""
if not self._warmup_end:
return
cur_epoch = runner.epoch
optimizer = runner.optim_wrapper.optimizer
for group_idx, param in enumerate(optimizer.param_groups):
param['lr'] = self._base_lr[group_idx] * self.scheduler_fn(
cur_epoch)
|