Threadpool fixes and CLI improvements (#550)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
This commit is contained in:
Glenn Jocher
2023-01-22 17:08:08 +01:00
committed by GitHub
parent d9a0fba251
commit 21b701c4ea
22 changed files with 338 additions and 251 deletions

View File

@ -126,15 +126,15 @@ class Exporter:
save_dir (Path): Directory to save results.
"""
def __init__(self, config=DEFAULT_CFG, overrides=None):
def __init__(self, cfg=DEFAULT_CFG, overrides=None):
"""
Initializes the Exporter class.
Args:
config (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
overrides (dict, optional): Configuration overrides. Defaults to None.
"""
self.args = get_cfg(config, overrides)
self.args = get_cfg(cfg, overrides)
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
callbacks.add_integration_callbacks(self)
@ -151,7 +151,7 @@ class Exporter:
# Load PyTorch model
self.device = select_device('cpu' if self.args.device is None else self.args.device)
if self.args.half:
if self.device.type == 'cpu' and not coreml:
if self.device.type == 'cpu' and not coreml and not xml:
LOGGER.info('half=True only compatible with GPU or CoreML export, i.e. use device=0 or format=coreml')
self.args.half = False
assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic'
@ -184,7 +184,7 @@ class Exporter:
y = None
for _ in range(2):
y = model(im) # dry runs
if self.args.half and not coreml:
if self.args.half and not coreml and not xml:
im, model = im.half(), model.half() # to FP16
shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape
LOGGER.info(
@ -332,7 +332,7 @@ class Exporter:
f = str(self.file).replace(self.file.suffix, f'_openvino_model{os.sep}')
f_onnx = self.file.with_suffix('.onnx')
cmd = f"mo --input_model {f_onnx} --output_dir {f} --data_type {'FP16' if self.args.half else 'FP32'}"
cmd = f"mo --input_model {f_onnx} --output_dir {f} {'--compress_to_fp16' * self.args.half}"
subprocess.run(cmd.split(), check=True, env=os.environ) # export
yaml_save(Path(f) / self.file.with_suffix('.yaml').name, self.metadata) # add metadata.yaml
return f, None

View File

@ -6,7 +6,7 @@ from ultralytics import yolo # noqa
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_one_weight
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.engine.exporter import Exporter
from ultralytics.yolo.utils import DEFAULT_CFG_PATH, LOGGER, yaml_load
from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, yaml_load
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.torch_utils import guess_task_from_head, smart_inference_mode
@ -151,7 +151,7 @@ class YOLO:
overrides = self.overrides.copy()
overrides.update(kwargs)
overrides["mode"] = "val"
args = get_cfg(cfg=DEFAULT_CFG_PATH, overrides=overrides)
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.data = data or args.data
args.task = self.task
@ -169,7 +169,7 @@ class YOLO:
overrides = self.overrides.copy()
overrides.update(kwargs)
args = get_cfg(cfg=DEFAULT_CFG_PATH, overrides=overrides)
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.task = self.task
print(args)
@ -181,8 +181,7 @@ class YOLO:
Trains the model on a given dataset.
Args:
**kwargs (Any): Any number of arguments representing the training configuration. List of all args can be found in 'config' section.
You can pass all arguments as a yaml file in `cfg`. Other args are ignored if `cfg` file is passed
**kwargs (Any): Any number of arguments representing the training configuration.
"""
overrides = self.overrides.copy()
overrides.update(kwargs)
@ -192,7 +191,7 @@ class YOLO:
overrides["task"] = self.task
overrides["mode"] = "train"
if not overrides.get("data"):
raise AttributeError("dataset not provided! Please define `data` in config.yaml or pass as an argument.")
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
if overrides.get("resume"):
overrides["resume"] = self.ckpt_path
@ -223,6 +222,13 @@ class YOLO:
return model_class, trainer_class, validator_class, predictor_class
@property
def names(self):
"""
Returns class names of the loaded model.
"""
return self.model.names
@staticmethod
def _reset_ckpt_args(args):
args.pop("project", None)

View File

@ -27,7 +27,6 @@ Usage - formats:
"""
import platform
from collections import defaultdict
from itertools import chain
from pathlib import Path
import cv2
@ -62,15 +61,15 @@ class BasePredictor:
data_path (str): Path to data.
"""
def __init__(self, config=DEFAULT_CFG_PATH, overrides=None):
def __init__(self, cfg=DEFAULT_CFG_PATH, overrides=None):
"""
Initializes the BasePredictor class.
Args:
config (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
overrides (dict, optional): Configuration overrides. Defaults to None.
"""
self.args = get_cfg(config, overrides)
self.args = get_cfg(cfg, overrides)
project = self.args.project or Path(SETTINGS['runs_dir']) / 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)
@ -219,7 +218,7 @@ class BasePredictor:
self.run_callbacks("on_predict_batch_end")
# Print results
if verbose:
if verbose and self.seen:
t = tuple(x.t / self.seen * 1E3 for x in self.dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms postprocess per image at shape '
f'{(1, 3, *self.imgsz)}' % t)

View File

@ -31,7 +31,8 @@ from ultralytics.yolo.utils.autobatch import check_train_batch_size
from ultralytics.yolo.utils.checks import check_file, check_imgsz, print_args
from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_command
from ultralytics.yolo.utils.files import get_latest_run, increment_path
from ultralytics.yolo.utils.torch_utils import ModelEMA, de_parallel, init_seeds, one_cycle, strip_optimizer
from ultralytics.yolo.utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, init_seeds, one_cycle,
strip_optimizer)
class BaseTrainer:
@ -71,15 +72,15 @@ class BaseTrainer:
csv (Path): Path to results CSV file.
"""
def __init__(self, config=DEFAULT_CFG_PATH, overrides=None):
def __init__(self, cfg=DEFAULT_CFG_PATH, overrides=None):
"""
Initializes the BaseTrainer class.
Args:
config (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CONFIG.
overrides (dict, optional): Configuration overrides. Defaults to None.
"""
self.args = get_cfg(config, overrides)
self.args = get_cfg(cfg, overrides)
self.device = utils.torch_utils.select_device(self.args.device, self.args.batch)
self.check_resume()
self.console = LOGGER
@ -225,6 +226,7 @@ class BaseTrainer:
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.last_epoch = self.start_epoch - 1 # do not move
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
@ -333,10 +335,12 @@ class BaseTrainer:
# Validation
self.ema.update_attr(self.model, include=['yaml', 'nc', 'args', 'names', 'stride', 'class_weights'])
final_epoch = (epoch + 1 == self.epochs)
final_epoch = (epoch + 1 == self.epochs) or self.stopper.possible_stop
if self.args.val or final_epoch:
self.metrics, self.fitness = self.validate()
self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **self.lr})
self.stop = self.stopper(epoch + 1, self.fitness)
# Save model
if self.args.save or (epoch + 1 == self.epochs):
@ -347,7 +351,15 @@ class BaseTrainer:
self.epoch_time = tnow - self.epoch_time_start
self.epoch_time_start = tnow
self.run_callbacks("on_fit_epoch_end")
# TODO: termination condition
# Early Stopping
if RANK != -1: # if DDP training
broadcast_list = [self.stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
self.stop = broadcast_list[0]
if self.stop:
break # must break all DDP ranks
if rank in {-1, 0}:
# Do final val with best.pt