HUB setup (#108)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia
2023-01-02 00:51:14 +05:30
committed by GitHub
parent c6eb6720de
commit 2bc9a5c87e
16 changed files with 631 additions and 122 deletions

View File

@ -3,46 +3,48 @@ from pathlib import Path
import hydra
import ultralytics
from ultralytics import yolo
from ultralytics import hub, yolo
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr
from .utils import DEFAULT_CONFIG, LOGGER, colorstr
DIR = Path(__file__).parent
@hydra.main(version_base=None, config_path="configs", config_name="default")
@hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent.relative_to(DIR)), config_name=DEFAULT_CONFIG.name)
def cli(cfg):
cwd = Path().cwd()
LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}")
"""
Run a specified task and mode with the given configuration.
Args:
cfg (DictConfig): Configuration for the task and mode.
"""
# LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}")
task, mode = cfg.task.lower(), cfg.mode.lower()
if task == "init": # special case
shutil.copy2(DEFAULT_CONFIG, cwd)
# Special case for initializing the configuration
if task == "init":
shutil.copy2(DEFAULT_CONFIG, Path.cwd())
LOGGER.info(f"""
{colorstr("YOLO:")} configuration saved to {cwd / DEFAULT_CONFIG.name}.
{colorstr("YOLO:")} configuration saved to {Path.cwd() / DEFAULT_CONFIG.name}.
To run experiments using custom configuration:
yolo task='task' mode='mode' --config-name config_file.yaml
""")
return
elif task == "detect":
module = yolo.v8.detect
elif task == "segment":
module = yolo.v8.segment
elif task == "classify":
module = yolo.v8.classify
elif task == "export":
func = yolo.engine.exporter.export
else:
raise SyntaxError("task not recognized. Choices are `'detect', 'segment', 'classify'`")
# Mapping from task to module
task_module_map = {"detect": yolo.v8.detect, "segment": yolo.v8.segment, "classify": yolo.v8.classify}
module = task_module_map.get(task)
if not module:
raise SyntaxError(f"task not recognized. Choices are {', '.join(task_module_map.keys())}")
# Mapping from mode to function
mode_func_map = {
"train": module.train,
"val": module.val,
"predict": module.predict,
"export": yolo.engine.exporter.export,
"checks": hub.checks}
func = mode_func_map.get(mode)
if not func:
raise SyntaxError(f"mode not recognized. Choices are {', '.join(mode_func_map.keys())}")
if mode == "train":
func = module.train
elif mode == "val":
func = module.val
elif mode == "predict":
func = module.predict
elif mode == "export":
func = yolo.engine.exporter.export
else:
raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict', 'export'`")
func(cfg)

View File

@ -8,6 +8,7 @@ mode: "train" # choices=['train', 'val', 'predict'] # mode to run task in.
model: null # i.e. yolov5s.pt, yolo.yaml. Path to model file
data: null # i.e. coco128.yaml. Path to data file
epochs: 100 # number of epochs to train for
patience: 50 # TODO: epochs to wait for no observable improvement for early stopping of training
batch_size: 16 # number of images per batch
imgsz: 640 # size of input images
save: True # save checkpoints

View File

@ -71,8 +71,7 @@ from ultralytics.nn.tasks import ClassificationModel, DetectionModel, Segmentati
from ultralytics.yolo.configs import get_config
from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages
from ultralytics.yolo.data.utils import check_dataset
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr, get_default_args, yaml_save
from ultralytics.yolo.utils.callbacks import default_callbacks
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, callbacks, colorstr, get_default_args, yaml_save
from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version, check_yaml
from ultralytics.yolo.utils.files import file_size, increment_path
from ultralytics.yolo.utils.ops import Profile
@ -138,16 +137,15 @@ class Exporter:
"""
if overrides is None:
overrides = {}
if 'batch_size' not in overrides:
overrides['batch_size'] = 1 # set default export batch size
self.args = get_config(config, overrides)
project = self.args.project or f"runs/{self.args.task}"
name = self.args.name or "exp" # hardcode mode as export doesn't require it
self.save_dir = increment_path(Path(project) / name, exist_ok=self.args.exist_ok)
self.save_dir.mkdir(parents=True, exist_ok=True)
# callbacks
self.callbacks = defaultdict([])
for callback, func in default_callbacks.items():
self.add_callback(callback, func)
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
callbacks.add_integration_callbacks(self)
@smart_inference_mode()
def __call__(self, model=None):
@ -173,7 +171,6 @@ class Exporter:
assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
# Input
self.args.batch_size = 1 # TODO: resolve this issue, default 16 not fit for export
im = torch.zeros(self.args.batch_size, 3, *self.imgsz).to(self.device)
file = Path(getattr(model, 'yaml_file', None) or Path(model.yaml['yaml_file']).name)
@ -765,18 +762,6 @@ class Exporter:
LOGGER.info(f'{prefix} pipeline success')
return model
def add_callback(self, event: str, callback):
"""
appends the given callback
"""
self.callbacks[event].append(callback)
def set_callback(self, event: str, callback):
"""
overrides the existing callbacks with the given callback
"""
self.callbacks[event] = [callback]
def run_callbacks(self, event: str):
for callback in self.callbacks.get(event, []):
callback(self)

View File

@ -35,8 +35,7 @@ from ultralytics.nn.autobackend import AutoBackend
from ultralytics.yolo.configs import get_config
from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages, LoadScreenshots, LoadStreams
from ultralytics.yolo.data.utils import IMG_FORMATS, VID_FORMATS
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr, ops
from ultralytics.yolo.utils.callbacks import default_callbacks
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, callbacks, colorstr, ops
from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_imshow
from ultralytics.yolo.utils.files import increment_path
from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode
@ -90,11 +89,8 @@ class BasePredictor:
self.view_img = None
self.annotator = None
self.data_path = None
# callbacks
self.callbacks = defaultdict([])
for callback, func in default_callbacks.items():
self.add_callback(callback, func)
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
callbacks.add_integration_callbacks(self)
def preprocess(self, img):
pass
@ -227,18 +223,6 @@ class BasePredictor:
self.vid_writer[idx] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
self.vid_writer[idx].write(im0)
def add_callback(self, event: str, callback):
"""
appends the given callback
"""
self.callbacks[event].append(callback)
def set_callback(self, event: str, callback):
"""
overrides the existing callbacks with the given callback
"""
self.callbacks[event] = [callback]
def run_callbacks(self, event: str):
for callback in self.callbacks.get(event, []):
callback(self)

View File

@ -21,11 +21,10 @@ from torch.optim import lr_scheduler
from tqdm import tqdm
import ultralytics.yolo.utils as utils
import ultralytics.yolo.utils.callbacks as callbacks
from ultralytics import __version__
from ultralytics.yolo.configs import get_config
from ultralytics.yolo.data.utils import check_dataset, check_dataset_yaml
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, RANK, TQDM_BAR_FORMAT, colorstr, yaml_save
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, RANK, TQDM_BAR_FORMAT, callbacks, colorstr, yaml_save
from ultralytics.yolo.utils.checks import check_file, print_args
from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_command
from ultralytics.yolo.utils.files import get_latest_run, increment_path
@ -88,7 +87,7 @@ class BaseTrainer:
self.model = None
self.callbacks = defaultdict(list)
# dirs
# Dirs
project = self.args.project or f"runs/{self.args.task}"
name = self.args.name or f"{self.args.mode}"
self.save_dir = increment_path(Path(project) / name, exist_ok=self.args.exist_ok if RANK in {-1, 0} else True)
@ -104,7 +103,7 @@ class BaseTrainer:
if RANK == -1:
print_args(dict(self.args))
# device
# Device
self.device = utils.torch_utils.select_device(self.args.device, self.batch_size)
self.amp = self.device.type != 'cpu'
self.scaler = amp.GradScaler(enabled=self.amp)
@ -123,7 +122,7 @@ class BaseTrainer:
self.lf = None
self.scheduler = None
# epoch level metrics
# Epoch level metrics
self.best_fitness = None
self.fitness = None
self.loss = None
@ -131,20 +130,20 @@ class BaseTrainer:
self.loss_names = None
self.csv = self.save_dir / 'results.csv'
for callback, func in callbacks.default_callbacks.items():
self.add_callback(callback, func)
# Callbacks
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
if RANK in {0, -1}:
callbacks.add_integration_callbacks(self)
def add_callback(self, event: str, callback):
"""
appends the given callback
Appends the given callback. TODO: unused, consider removing
"""
self.callbacks[event].append(callback)
def set_callback(self, event: str, callback):
"""
overrides the existing callbacks with the given callback
Overrides the existing callbacks with the given callback. TODO: unused, consider removing
"""
self.callbacks[event] = [callback]
@ -469,7 +468,7 @@ class BaseTrainer:
self.validator.args.save_json = True
self.metrics = self.validator(model=f)
self.metrics.pop('fitness', None)
self.run_callbacks('on_val_end')
self.run_callbacks('on_fit_epoch_end')
def check_resume(self):
resume = self.args.resume

View File

@ -8,8 +8,7 @@ from tqdm import tqdm
from ultralytics.nn.autobackend import AutoBackend
from ultralytics.yolo.data.utils import check_dataset, check_dataset_yaml
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, RANK, TQDM_BAR_FORMAT
from ultralytics.yolo.utils.callbacks import default_callbacks
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, RANK, TQDM_BAR_FORMAT, callbacks
from ultralytics.yolo.utils.checks import check_imgsz
from ultralytics.yolo.utils.files import increment_path
from ultralytics.yolo.utils.ops import Profile
@ -66,10 +65,7 @@ class BaseValidator:
exist_ok=self.args.exist_ok if RANK in {-1, 0} else True)
(self.save_dir / 'labels' if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True)
# callbacks
self.callbacks = defaultdict(list)
for callback, func in default_callbacks.items():
self.add_callback(callback, func)
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
@smart_inference_mode()
def __call__(self, trainer=None, model=None):
@ -77,7 +73,6 @@ class BaseValidator:
Supports validation of a pre-trained model if passed or a model being trained
if trainer is passed (trainer gets priority).
"""
self.run_callbacks('on_val_start')
self.training = trainer is not None
if self.training:
self.device = trainer.device
@ -89,6 +84,8 @@ class BaseValidator:
self.loss = torch.zeros_like(trainer.loss_items, device=trainer.device)
self.args.plots = trainer.epoch == trainer.epochs - 1 # always plot final epoch
else:
callbacks.add_integration_callbacks(self)
self.run_callbacks('on_val_start')
assert model is not None, "Either trainer or model is needed for validation"
self.device = select_device(self.args.device, self.args.batch_size)
self.args.half &= self.device.type != 'cpu'
@ -167,18 +164,6 @@ class BaseValidator:
stats = self.eval_json(stats) # update stats
return stats
def add_callback(self, event: str, callback):
"""
appends the given callback
"""
self.callbacks[event].append(callback)
def set_callback(self, event: str, callback):
"""
overrides the existing callbacks with the given callback
"""
self.callbacks[event] = [callback]
def run_callbacks(self, event: str):
for callback in self.callbacks.get(event, []):
callback(self)

View File

@ -249,26 +249,6 @@ def threaded(func):
return wrapper
def get_settings(file=USER_CONFIG_DIR / 'settings.yaml'):
"""
Function that loads a global settings YAML, or creates it and populates it with default values if it does not exist.
If the datasets or weights directories are set to None, the current working directory will be used.
The 'sync' setting determines whether analytics will be synced to help with YOLO development.
"""
from ultralytics.yolo.utils.torch_utils import torch_distributed_zero_first
with torch_distributed_zero_first(RANK):
if not file.exists():
settings = {
'datasets_dir': None, # default datasets directory. If None, current working directory is used.
'weights_dir': None, # default weights directory. If None, current working directory is used.
'sync': True} # sync analytics to help with YOLO development
yaml_save(file, settings)
return yaml_load(file)
def yaml_save(file='data.yaml', data=None):
"""
Save YAML data to a file.
@ -305,6 +285,26 @@ def yaml_load(file='data.yaml'):
return {**yaml.safe_load(f), 'yaml_file': file}
def get_settings(file=USER_CONFIG_DIR / 'settings.yaml'):
"""
Function that loads a global settings YAML, or creates it and populates it with default values if it does not exist.
If the datasets or weights directories are set to None, the current working directory will be used.
The 'sync' setting determines whether analytics will be synced to help with YOLO development.
"""
from ultralytics.yolo.utils.torch_utils import torch_distributed_zero_first
with torch_distributed_zero_first(RANK):
if not file.exists():
settings = {
'datasets_dir': None, # default datasets directory. If None, current working directory is used.
'weights_dir': None, # default weights directory. If None, current working directory is used.
'sync': True} # sync analytics to help with YOLO development
yaml_save(file, settings)
return yaml_load(file)
# Run below code on utils init -----------------------------------------------------------------------------------------
# Set logger

View File

@ -135,11 +135,12 @@ default_callbacks = {
'on_export_end': on_export_end}
def add_integration_callbacks(trainer):
def add_integration_callbacks(instance):
from .clearml import callbacks as clearml_callbacks
from .tb import callbacks as tb_callbacks
from .hub import callbacks as hub_callbacks
from .tensorboard import callbacks as tb_callbacks
from .wb import callbacks as wb_callbacks
for x in clearml_callbacks, tb_callbacks, wb_callbacks:
for x in clearml_callbacks, hub_callbacks, tb_callbacks, wb_callbacks:
for k, v in x.items():
trainer.add_callback(k, v) # add_callback(name, func)
instance.callbacks[k].append(v) # callback[name].append(func)

View File

@ -0,0 +1,80 @@
import json
from time import time
import torch
from ultralytics.hub.utils import PREFIX, sync_analytics
from ultralytics.yolo.utils import LOGGER
def on_pretrain_routine_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Start timer for upload rate limit
LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
session.t = {'metrics': time(), 'ckpt': time()} # start timer on self.rate_limit
def on_fit_epoch_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Upload metrics after val end
metrics = trainer.metrics
for k, v in metrics.items():
if isinstance(v, torch.Tensor):
metrics[k] = v.item()
session.metrics_queue[trainer.epoch] = json.dumps(metrics) # json string
if time() - session.t['metrics'] > session.rate_limits['metrics']:
session.upload_metrics()
session.t['metrics'] = time() # reset timer
session.metrics_queue = {} # reset queue
def on_model_save(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Upload checkpoints with rate limiting
is_best = trainer.best_fitness == trainer.fitness
if time() - session.t['ckpt'] > session.rate_limits['ckpt']:
LOGGER.info(f"{PREFIX}Uploading checkpoint {session.model_id}")
session.upload_model(trainer.epoch, trainer.last, is_best)
session.t['ckpt'] = time() # reset timer
def on_train_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Upload final model and metrics with exponential standoff
LOGGER.info(f"{PREFIX}Training completed successfully ✅\n"
f"{PREFIX}Uploading final {session.model_id}")
session.upload_model(trainer.epoch, trainer.best, map=trainer.metrics['metrics/mAP50(B)'], final=True)
session.alive = False # stop heartbeats
LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
def on_train_start(trainer):
sync_analytics(trainer.args)
def on_val_start(validator):
sync_analytics(validator.args)
def on_predict_start(predictor):
sync_analytics(predictor.args)
def on_export_start(exporter):
sync_analytics(exporter.args)
callbacks = {
"on_pretrain_routine_end": on_pretrain_routine_end,
"on_fit_epoch_end": on_fit_epoch_end,
"on_model_save": on_model_save,
"on_train_end": on_train_end,
"on_train_start": on_train_start,
"on_val_start": on_val_start,
"on_predict_start": on_predict_start,
"on_export_start": on_export_start}