From 451cf8b6473ebd3363f74e5321203889ec9c6261 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 4 Jun 2023 22:35:50 +0200 Subject: [PATCH] Add Adamax, NAdam, RAdam optimizers (#2969) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/modes/train.md | 94 +++++++++++++++--------------- docs/usage/cfg.md | 94 +++++++++++++++--------------- ultralytics/yolo/cfg/default.yaml | 2 +- ultralytics/yolo/engine/trainer.py | 80 ++++++++++++++----------- ultralytics/yolo/utils/tuner.py | 2 +- 5 files changed, 141 insertions(+), 131 deletions(-) diff --git a/docs/modes/train.md b/docs/modes/train.md index d560360..4492373 100644 --- a/docs/modes/train.md +++ b/docs/modes/train.md @@ -55,50 +55,50 @@ include the choice of optimizer, the choice of loss function, and the size and c is important to carefully tune and experiment with these settings to achieve the best possible performance for a given task. -| Key | Value | Description | -|-------------------|----------|-----------------------------------------------------------------------------| -| `model` | `None` | path to model file, i.e. yolov8n.pt, yolov8n.yaml | -| `data` | `None` | path to data file, i.e. coco128.yaml | -| `epochs` | `100` | number of epochs to train for | -| `patience` | `50` | epochs to wait for no observable improvement for early stopping of training | -| `batch` | `16` | number of images per batch (-1 for AutoBatch) | -| `imgsz` | `640` | size of input images as integer or w,h | -| `save` | `True` | save train checkpoints and predict results | -| `save_period` | `-1` | Save checkpoint every x epochs (disabled if < 1) | -| `cache` | `False` | True/ram, disk or False. Use cache for data loading | -| `device` | `None` | device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu | -| `workers` | `8` | number of worker threads for data loading (per RANK if DDP) | -| `project` | `None` | project name | -| `name` | `None` | experiment name | -| `exist_ok` | `False` | whether to overwrite existing experiment | -| `pretrained` | `False` | whether to use a pretrained model | -| `optimizer` | `'SGD'` | optimizer to use, choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] | -| `verbose` | `False` | whether to print verbose output | -| `seed` | `0` | random seed for reproducibility | -| `deterministic` | `True` | whether to enable deterministic mode | -| `single_cls` | `False` | train multi-class data as single-class | -| `rect` | `False` | rectangular training with each batch collated for minimum padding | -| `cos_lr` | `False` | use cosine learning rate scheduler | -| `close_mosaic` | `0` | (int) disable mosaic augmentation for final epochs | -| `resume` | `False` | resume training from last checkpoint | -| `amp` | `True` | Automatic Mixed Precision (AMP) training, choices=[True, False] | -| `fraction` | `1.0` | dataset fraction to train on (default is 1.0, all images in train set) | -| `profile` | `False` | profile ONNX and TensorRT speeds during training for loggers | -| `lr0` | `0.01` | initial learning rate (i.e. SGD=1E-2, Adam=1E-3) | -| `lrf` | `0.01` | final learning rate (lr0 * lrf) | -| `momentum` | `0.937` | SGD momentum/Adam beta1 | -| `weight_decay` | `0.0005` | optimizer weight decay 5e-4 | -| `warmup_epochs` | `3.0` | warmup epochs (fractions ok) | -| `warmup_momentum` | `0.8` | warmup initial momentum | -| `warmup_bias_lr` | `0.1` | warmup initial bias lr | -| `box` | `7.5` | box loss gain | -| `cls` | `0.5` | cls loss gain (scale with pixels) | -| `dfl` | `1.5` | dfl loss gain | -| `pose` | `12.0` | pose loss gain (pose-only) | -| `kobj` | `2.0` | keypoint obj loss gain (pose-only) | -| `label_smoothing` | `0.0` | label smoothing (fraction) | -| `nbs` | `64` | nominal batch size | -| `overlap_mask` | `True` | masks should overlap during training (segment train only) | -| `mask_ratio` | `4` | mask downsample ratio (segment train only) | -| `dropout` | `0.0` | use dropout regularization (classify train only) | -| `val` | `True` | validate/test during training | +| Key | Value | Description | +|-------------------|----------|-----------------------------------------------------------------------------------| +| `model` | `None` | path to model file, i.e. yolov8n.pt, yolov8n.yaml | +| `data` | `None` | path to data file, i.e. coco128.yaml | +| `epochs` | `100` | number of epochs to train for | +| `patience` | `50` | epochs to wait for no observable improvement for early stopping of training | +| `batch` | `16` | number of images per batch (-1 for AutoBatch) | +| `imgsz` | `640` | size of input images as integer or w,h | +| `save` | `True` | save train checkpoints and predict results | +| `save_period` | `-1` | Save checkpoint every x epochs (disabled if < 1) | +| `cache` | `False` | True/ram, disk or False. Use cache for data loading | +| `device` | `None` | device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu | +| `workers` | `8` | number of worker threads for data loading (per RANK if DDP) | +| `project` | `None` | project name | +| `name` | `None` | experiment name | +| `exist_ok` | `False` | whether to overwrite existing experiment | +| `pretrained` | `False` | whether to use a pretrained model | +| `optimizer` | `'auto'` | optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto] | +| `verbose` | `False` | whether to print verbose output | +| `seed` | `0` | random seed for reproducibility | +| `deterministic` | `True` | whether to enable deterministic mode | +| `single_cls` | `False` | train multi-class data as single-class | +| `rect` | `False` | rectangular training with each batch collated for minimum padding | +| `cos_lr` | `False` | use cosine learning rate scheduler | +| `close_mosaic` | `0` | (int) disable mosaic augmentation for final epochs | +| `resume` | `False` | resume training from last checkpoint | +| `amp` | `True` | Automatic Mixed Precision (AMP) training, choices=[True, False] | +| `fraction` | `1.0` | dataset fraction to train on (default is 1.0, all images in train set) | +| `profile` | `False` | profile ONNX and TensorRT speeds during training for loggers | +| `lr0` | `0.01` | initial learning rate (i.e. SGD=1E-2, Adam=1E-3) | +| `lrf` | `0.01` | final learning rate (lr0 * lrf) | +| `momentum` | `0.937` | SGD momentum/Adam beta1 | +| `weight_decay` | `0.0005` | optimizer weight decay 5e-4 | +| `warmup_epochs` | `3.0` | warmup epochs (fractions ok) | +| `warmup_momentum` | `0.8` | warmup initial momentum | +| `warmup_bias_lr` | `0.1` | warmup initial bias lr | +| `box` | `7.5` | box loss gain | +| `cls` | `0.5` | cls loss gain (scale with pixels) | +| `dfl` | `1.5` | dfl loss gain | +| `pose` | `12.0` | pose loss gain (pose-only) | +| `kobj` | `2.0` | keypoint obj loss gain (pose-only) | +| `label_smoothing` | `0.0` | label smoothing (fraction) | +| `nbs` | `64` | nominal batch size | +| `overlap_mask` | `True` | masks should overlap during training (segment train only) | +| `mask_ratio` | `4` | mask downsample ratio (segment train only) | +| `dropout` | `0.0` | use dropout regularization (classify train only) | +| `val` | `True` | validate/test during training | diff --git a/docs/usage/cfg.md b/docs/usage/cfg.md index 0c99002..b7ddb49 100644 --- a/docs/usage/cfg.md +++ b/docs/usage/cfg.md @@ -77,53 +77,53 @@ include: The training settings for YOLO models encompass various hyperparameters and configurations used during the training process. These settings influence the model's performance, speed, and accuracy. Key training settings include batch size, learning rate, momentum, and weight decay. Additionally, the choice of optimizer, loss function, and training dataset composition can impact the training process. Careful tuning and experimentation with these settings are crucial for optimizing performance. -| Key | Value | Description | -|-------------------|----------|-----------------------------------------------------------------------------| -| `model` | `None` | path to model file, i.e. yolov8n.pt, yolov8n.yaml | -| `data` | `None` | path to data file, i.e. coco128.yaml | -| `epochs` | `100` | number of epochs to train for | -| `patience` | `50` | epochs to wait for no observable improvement for early stopping of training | -| `batch` | `16` | number of images per batch (-1 for AutoBatch) | -| `imgsz` | `640` | size of input images as integer or w,h | -| `save` | `True` | save train checkpoints and predict results | -| `save_period` | `-1` | Save checkpoint every x epochs (disabled if < 1) | -| `cache` | `False` | True/ram, disk or False. Use cache for data loading | -| `device` | `None` | device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu | -| `workers` | `8` | number of worker threads for data loading (per RANK if DDP) | -| `project` | `None` | project name | -| `name` | `None` | experiment name | -| `exist_ok` | `False` | whether to overwrite existing experiment | -| `pretrained` | `False` | whether to use a pretrained model | -| `optimizer` | `'SGD'` | optimizer to use, choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] | -| `verbose` | `False` | whether to print verbose output | -| `seed` | `0` | random seed for reproducibility | -| `deterministic` | `True` | whether to enable deterministic mode | -| `single_cls` | `False` | train multi-class data as single-class | -| `rect` | `False` | rectangular training with each batch collated for minimum padding | -| `cos_lr` | `False` | use cosine learning rate scheduler | -| `close_mosaic` | `0` | (int) disable mosaic augmentation for final epochs | -| `resume` | `False` | resume training from last checkpoint | -| `amp` | `True` | Automatic Mixed Precision (AMP) training, choices=[True, False] | -| `fraction` | `1.0` | dataset fraction to train on (default is 1.0, all images in train set) | -| `profile` | `False` | profile ONNX and TensorRT speeds during training for loggers | -| `lr0` | `0.01` | initial learning rate (i.e. SGD=1E-2, Adam=1E-3) | -| `lrf` | `0.01` | final learning rate (lr0 * lrf) | -| `momentum` | `0.937` | SGD momentum/Adam beta1 | -| `weight_decay` | `0.0005` | optimizer weight decay 5e-4 | -| `warmup_epochs` | `3.0` | warmup epochs (fractions ok) | -| `warmup_momentum` | `0.8` | warmup initial momentum | -| `warmup_bias_lr` | `0.1` | warmup initial bias lr | -| `box` | `7.5` | box loss gain | -| `cls` | `0.5` | cls loss gain (scale with pixels) | -| `dfl` | `1.5` | dfl loss gain | -| `pose` | `12.0` | pose loss gain (pose-only) | -| `kobj` | `2.0` | keypoint obj loss gain (pose-only) | -| `label_smoothing` | `0.0` | label smoothing (fraction) | -| `nbs` | `64` | nominal batch size | -| `overlap_mask` | `True` | masks should overlap during training (segment train only) | -| `mask_ratio` | `4` | mask downsample ratio (segment train only) | -| `dropout` | `0.0` | use dropout regularization (classify train only) | -| `val` | `True` | validate/test during training | +| Key | Value | Description | +|-------------------|----------|-----------------------------------------------------------------------------------| +| `model` | `None` | path to model file, i.e. yolov8n.pt, yolov8n.yaml | +| `data` | `None` | path to data file, i.e. coco128.yaml | +| `epochs` | `100` | number of epochs to train for | +| `patience` | `50` | epochs to wait for no observable improvement for early stopping of training | +| `batch` | `16` | number of images per batch (-1 for AutoBatch) | +| `imgsz` | `640` | size of input images as integer or w,h | +| `save` | `True` | save train checkpoints and predict results | +| `save_period` | `-1` | Save checkpoint every x epochs (disabled if < 1) | +| `cache` | `False` | True/ram, disk or False. Use cache for data loading | +| `device` | `None` | device to run on, i.e. cuda device=0 or device=0,1,2,3 or device=cpu | +| `workers` | `8` | number of worker threads for data loading (per RANK if DDP) | +| `project` | `None` | project name | +| `name` | `None` | experiment name | +| `exist_ok` | `False` | whether to overwrite existing experiment | +| `pretrained` | `False` | whether to use a pretrained model | +| `optimizer` | `'auto'` | optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto] | +| `verbose` | `False` | whether to print verbose output | +| `seed` | `0` | random seed for reproducibility | +| `deterministic` | `True` | whether to enable deterministic mode | +| `single_cls` | `False` | train multi-class data as single-class | +| `rect` | `False` | rectangular training with each batch collated for minimum padding | +| `cos_lr` | `False` | use cosine learning rate scheduler | +| `close_mosaic` | `0` | (int) disable mosaic augmentation for final epochs | +| `resume` | `False` | resume training from last checkpoint | +| `amp` | `True` | Automatic Mixed Precision (AMP) training, choices=[True, False] | +| `fraction` | `1.0` | dataset fraction to train on (default is 1.0, all images in train set) | +| `profile` | `False` | profile ONNX and TensorRT speeds during training for loggers | +| `lr0` | `0.01` | initial learning rate (i.e. SGD=1E-2, Adam=1E-3) | +| `lrf` | `0.01` | final learning rate (lr0 * lrf) | +| `momentum` | `0.937` | SGD momentum/Adam beta1 | +| `weight_decay` | `0.0005` | optimizer weight decay 5e-4 | +| `warmup_epochs` | `3.0` | warmup epochs (fractions ok) | +| `warmup_momentum` | `0.8` | warmup initial momentum | +| `warmup_bias_lr` | `0.1` | warmup initial bias lr | +| `box` | `7.5` | box loss gain | +| `cls` | `0.5` | cls loss gain (scale with pixels) | +| `dfl` | `1.5` | dfl loss gain | +| `pose` | `12.0` | pose loss gain (pose-only) | +| `kobj` | `2.0` | keypoint obj loss gain (pose-only) | +| `label_smoothing` | `0.0` | label smoothing (fraction) | +| `nbs` | `64` | nominal batch size | +| `overlap_mask` | `True` | masks should overlap during training (segment train only) | +| `mask_ratio` | `4` | mask downsample ratio (segment train only) | +| `dropout` | `0.0` | use dropout regularization (classify train only) | +| `val` | `True` | validate/test during training | [Train Guide](../modes/train.md){ .md-button .md-button--primary} diff --git a/ultralytics/yolo/cfg/default.yaml b/ultralytics/yolo/cfg/default.yaml index 41b9449..35be2e9 100644 --- a/ultralytics/yolo/cfg/default.yaml +++ b/ultralytics/yolo/cfg/default.yaml @@ -20,7 +20,7 @@ project: # project name name: # experiment name, results saved to 'project/name' directory exist_ok: False # whether to overwrite existing experiment pretrained: False # whether to use a pretrained model -optimizer: SGD # optimizer to use, choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] +optimizer: auto # optimizer to use, choices=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto] verbose: True # whether to print verbose output seed: 0 # random seed for reproducibility deterministic: True # whether to enable deterministic mode diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index b39ad13..b0da4df 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -5,6 +5,7 @@ Train a model on a dataset Usage: $ yolo mode=train model=yolov8n.pt data=coco128.yaml imgsz=640 epochs=100 batch=16 """ +import math import os import subprocess import time @@ -14,11 +15,10 @@ from pathlib import Path import numpy as np import torch -import torch.distributed as dist -import torch.nn as nn +from torch import distributed as dist +from torch import nn, optim from torch.cuda import amp from torch.nn.parallel import DistributedDataParallel as DDP -from torch.optim import lr_scheduler from tqdm import tqdm from ultralytics.nn.tasks import attempt_load_one_weight, attempt_load_weights @@ -234,33 +234,35 @@ class BaseTrainer: SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. ' 'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16') + # Dataloaders + batch_size = self.batch_size // max(world_size, 1) + self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode='train') + if RANK in (-1, 0): + self.test_loader = self.get_dataloader(self.testset, batch_size=batch_size * 2, rank=-1, mode='val') + self.validator = self.get_validator() + metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix='val') + self.metrics = dict(zip(metric_keys, [0] * len(metric_keys))) # TODO: init metrics for plot_results()? + self.ema = ModelEMA(self.model) + if self.args.plots and not self.args.v5loader: + self.plot_training_labels() + # Optimizer self.accumulate = max(round(self.args.nbs / self.batch_size), 1) # accumulate loss before optimizing weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs # scale weight_decay + iterations = math.ceil(len(self.train_loader.dataset) / max(self.batch_size, self.args.nbs)) * self.epochs self.optimizer = self.build_optimizer(model=self.model, name=self.args.optimizer, lr=self.args.lr0, momentum=self.args.momentum, - decay=weight_decay) + decay=weight_decay, + iterations=iterations) # Scheduler if self.args.cos_lr: self.lf = one_cycle(1, self.args.lrf, self.epochs) # cosine 1->hyp['lrf'] else: self.lf = lambda x: (1 - x / self.epochs) * (1.0 - self.args.lrf) + self.args.lrf # linear - self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf) + self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf) self.stopper, self.stop = EarlyStopping(patience=self.args.patience), False - - # Dataloaders - batch_size = self.batch_size // world_size if world_size > 1 else self.batch_size - self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode='train') - if RANK in (-1, 0): - self.test_loader = self.get_dataloader(self.testset, batch_size=batch_size * 2, rank=-1, mode='val') - self.validator = self.get_validator() - metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix='val') - self.metrics = dict(zip(metric_keys, [0] * len(metric_keys))) # TODO: init metrics for plot_results()? - self.ema = ModelEMA(self.model) - if self.args.plots and not self.args.v5loader: - self.plot_training_labels() self.resume_training(ckpt) self.scheduler.last_epoch = self.start_epoch - 1 # do not move self.run_callbacks('on_pretrain_routine_end') @@ -603,24 +605,30 @@ class BaseTrainer: if hasattr(self.train_loader.dataset, 'close_mosaic'): self.train_loader.dataset.close_mosaic(hyp=self.args) - @staticmethod - def build_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5): + def build_optimizer(self, model, name='auto', lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5): """ - Builds an optimizer with the specified parameters and parameter groups. + Constructs an optimizer for the given model, based on the specified optimizer name, learning rate, + momentum, weight decay, and number of iterations. Args: - model (nn.Module): model to optimize - name (str): name of the optimizer to use - lr (float): learning rate - momentum (float): momentum - decay (float): weight decay + model (torch.nn.Module): The model for which to build an optimizer. + name (str, optional): The name of the optimizer to use. If 'auto', the optimizer is selected + based on the number of iterations. Default: 'auto'. + lr (float, optional): The learning rate for the optimizer. Default: 0.001. + momentum (float, optional): The momentum factor for the optimizer. Default: 0.9. + decay (float, optional): The weight decay for the optimizer. Default: 1e-5. + iterations (float, optional): The number of iterations, which determines the optimizer if + name is 'auto'. Default: 1e5. Returns: - optimizer (torch.optim.Optimizer): the built optimizer + (torch.optim.Optimizer): The constructed optimizer. """ g = [], [], [] # optimizer parameter groups bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k) # normalization layers, i.e. BatchNorm2d() + if name == 'auto': + name, lr, momentum = ('SGD', 0.01, 0.9) if iterations > 6000 else ('NAdam', 0.001, 0.9) + self.args.warmup_bias_lr = 0.0 # no higher than 0.01 for NAdam for module_name, module in model.named_modules(): for param_name, param in module.named_parameters(recurse=False): @@ -632,19 +640,21 @@ class BaseTrainer: else: # weight (with decay) g[0].append(param) - if name == 'Adam': - optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999)) # adjust beta1 to momentum - elif name == 'AdamW': - optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0) + if name in ('Adam', 'Adamax', 'AdamW', 'NAdam', 'RAdam'): + optimizer = getattr(optim, name, optim.Adam)(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0) elif name == 'RMSProp': - optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum) + optimizer = optim.RMSprop(g[2], lr=lr, momentum=momentum) elif name == 'SGD': - optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True) + optimizer = optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True) else: - raise NotImplementedError(f'Optimizer {name} not implemented.') + raise NotImplementedError( + f"Optimizer '{name}' not found in list of available optimizers " + f'[Adam, AdamW, NAdam, RAdam, RMSProp, SGD, auto].' + 'To request support for addition optimizers please visit https://github.com/ultralytics/ultralytics.') optimizer.add_param_group({'params': g[0], 'weight_decay': decay}) # add g0 with weight_decay optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0}) # add g1 (BatchNorm2d weights) - LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups " - f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias') + LOGGER.info( + f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}, momentum={momentum}) with parameter groups " + f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias(decay=0.0)') return optimizer diff --git a/ultralytics/yolo/utils/tuner.py b/ultralytics/yolo/utils/tuner.py index 54e1b01..9f57677 100644 --- a/ultralytics/yolo/utils/tuner.py +++ b/ultralytics/yolo/utils/tuner.py @@ -14,7 +14,7 @@ except ImportError: tune = None default_space = { - # 'optimizer': tune.choice(['SGD', 'Adam', 'AdamW', 'RMSProp']), + # 'optimizer': tune.choice(['SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp']), 'lr0': tune.uniform(1e-5, 1e-1), 'lrf': tune.uniform(0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) 'momentum': tune.uniform(0.6, 0.98), # SGD momentum/Adam beta1