lichen/ieemoo-ai-contrast
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

增加学习率调度方式

Browse Source
This commit is contained in:
lee
2025-06-13 10:45:53 +08:00
parent 37ecef40f7
commit 1803f319a5
13 changed files with 319 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
configs/utils.py
View File

@ -1,4 +1,4 @@
from model import (resnet18, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
from model import (resnet18, resnet34, resnet50, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
PPLCNET_x1_0, PPLCNET_x0_5, PPLCNET_x2_5)
from timm.models import vit_base_patch16_224 as vit_base_16
from model.metric import ArcFace, CosFace
@ -14,6 +14,8 @@ class trainer_tools:
def get_backbone(self):
backbone_mapping = {
'resnet18': lambda: resnet18(scale=self.conf['models']['channel_ratio']),
'resnet34': lambda: resnet34(scale=self.conf['models']['channel_ratio']),
'resnet50': lambda: resnet50(scale=self.conf['models']['channel_ratio']),
'mobilevit_s': lambda: mobilevit_s(),
'mobilenetv3_small': lambda: MobileNetV3_Small(),
'PPLCNET_x1_0': lambda: PPLCNET_x1_0(),
@ -54,3 +56,24 @@ class trainer_tools:
)
}
return optimizer_mapping
def get_scheduler(self, optimizer):
scheduler_mapping = {
'step': lambda: optim.lr_scheduler.StepLR(
optimizer,
step_size=self.conf['training']['lr_step'],
gamma=self.conf['training']['lr_decay']
),
'cosine': lambda: optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=self.conf['training']['epochs'],
eta_min=self.conf['training']['cosine_eta_min']
),
'cosine_warm': lambda: optim.lr_scheduler.CosineAnnealingWarmRestarts(
optimizer,
T_0=self.conf['training'].get('cosine_t_0', 10),
T_mult=self.conf['training'].get('cosine_t_mult', 1),
eta_min=self.conf['training'].get('cosine_eta_min', 0)
)
}
return scheduler_mapping