ultralytics 8.0.81 single-line docstring updates (#2061)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
@ -182,8 +182,10 @@ def plt_settings(rcparams={'font.size': 11}, backend='Agg'):
|
||||
"""
|
||||
|
||||
def decorator(func):
|
||||
"""Decorator to apply temporary rc parameters and backend to a function."""
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
"""Sets rc parameters and backend, calls the original function, and restores the settings."""
|
||||
original_backend = plt.get_backend()
|
||||
plt.switch_backend(backend)
|
||||
|
||||
@ -229,6 +231,7 @@ class EmojiFilter(logging.Filter):
|
||||
"""
|
||||
|
||||
def filter(self, record):
|
||||
"""Filter logs by emoji unicode characters on windows."""
|
||||
record.msg = emojis(record.msg)
|
||||
return super().filter(record)
|
||||
|
||||
@ -573,13 +576,16 @@ class TryExcept(contextlib.ContextDecorator):
|
||||
"""YOLOv8 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager."""
|
||||
|
||||
def __init__(self, msg='', verbose=True):
|
||||
"""Initialize TryExcept class with optional message and verbosity settings."""
|
||||
self.msg = msg
|
||||
self.verbose = verbose
|
||||
|
||||
def __enter__(self):
|
||||
"""Executes when entering TryExcept context, initializes instance."""
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type, value, traceback):
|
||||
"""Defines behavior when exiting a 'with' block, prints error message if necessary."""
|
||||
if self.verbose and value:
|
||||
print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}"))
|
||||
return True
|
||||
@ -589,6 +595,7 @@ def threaded(func):
|
||||
"""Multi-threads a target function and returns thread. Usage: @threaded decorator."""
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
"""Multi-threads a given function and returns the thread."""
|
||||
thread = threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
@ -602,6 +609,7 @@ def set_sentry():
|
||||
"""
|
||||
|
||||
def before_send(event, hint):
|
||||
"""A function executed before sending the event to Sentry."""
|
||||
if 'exc_info' in hint:
|
||||
exc_type, exc_value, tb = hint['exc_info']
|
||||
if exc_type in (KeyboardInterrupt, FileNotFoundError) \
|
||||
@ -698,6 +706,7 @@ def set_settings(kwargs, file=SETTINGS_YAML):
|
||||
|
||||
|
||||
def deprecation_warn(arg, new_arg, version=None):
|
||||
"""Issue a deprecation warning when a deprecated argument is used, suggesting an updated argument."""
|
||||
if not version:
|
||||
version = float(__version__[:3]) + 0.2 # deprecate after 2nd major release
|
||||
LOGGER.warning(f"WARNING ⚠️ '{arg}' is deprecated and will be removed in 'ultralytics {version}' in the future. "
|
||||
|
||||
@ -35,7 +35,30 @@ from ultralytics.yolo.utils.files import file_size
|
||||
from ultralytics.yolo.utils.torch_utils import select_device
|
||||
|
||||
|
||||
def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', imgsz=160, half=False, device='cpu', hard_fail=False):
|
||||
def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt',
|
||||
imgsz=160,
|
||||
half=False,
|
||||
int8=False,
|
||||
device='cpu',
|
||||
hard_fail=False):
|
||||
"""
|
||||
Benchmark a YOLO model across different formats for speed and accuracy.
|
||||
|
||||
Args:
|
||||
model (Union[str, Path], optional): Path to the model file or directory. Default is
|
||||
Path(SETTINGS['weights_dir']) / 'yolov8n.pt'.
|
||||
imgsz (int, optional): Image size for the benchmark. Default is 160.
|
||||
half (bool, optional): Use half-precision for the model if True. Default is False.
|
||||
int8 (bool, optional): Use int8-precision for the model if True. Default is False.
|
||||
device (str, optional): Device to run the benchmark on, either 'cpu' or 'cuda'. Default is 'cpu'.
|
||||
hard_fail (Union[bool, float], optional): If True or a float, assert benchmarks pass with given metric.
|
||||
Default is False.
|
||||
|
||||
Returns:
|
||||
df (pandas.DataFrame): A pandas DataFrame with benchmark results for each format, including file size,
|
||||
metric, and inference time.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
pd.options.display.max_columns = 10
|
||||
pd.options.display.width = 120
|
||||
@ -61,7 +84,7 @@ def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', imgsz=160, hal
|
||||
filename = model.ckpt_path or model.cfg
|
||||
export = model # PyTorch format
|
||||
else:
|
||||
filename = model.export(imgsz=imgsz, format=format, half=half, device=device) # all others
|
||||
filename = model.export(imgsz=imgsz, format=format, half=half, int8=int8, device=device) # all others
|
||||
export = YOLO(filename, task=model.task)
|
||||
assert suffix in str(filename), 'export failed'
|
||||
emoji = '❎' # indicates export succeeded
|
||||
@ -83,7 +106,14 @@ def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', imgsz=160, hal
|
||||
elif model.task == 'pose':
|
||||
data, key = 'coco8-pose.yaml', 'metrics/mAP50-95(P)'
|
||||
|
||||
results = export.val(data=data, batch=1, imgsz=imgsz, plots=False, device=device, half=half, verbose=False)
|
||||
results = export.val(data=data,
|
||||
batch=1,
|
||||
imgsz=imgsz,
|
||||
plots=False,
|
||||
device=device,
|
||||
half=half,
|
||||
int8=int8,
|
||||
verbose=False)
|
||||
metric, speed = results.results_dict[key], results.speed['inference']
|
||||
y.append([name, '✅', round(file_size(filename), 1), round(metric, 4), round(speed, 2)])
|
||||
except Exception as e:
|
||||
|
||||
@ -2,111 +2,144 @@
|
||||
"""
|
||||
Base callbacks
|
||||
"""
|
||||
|
||||
from collections import defaultdict
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
# Trainer callbacks ----------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Called before the pretraining routine starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Called after the pretraining routine ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_start(trainer):
|
||||
"""Called when the training starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_epoch_start(trainer):
|
||||
"""Called at the start of each training epoch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_batch_start(trainer):
|
||||
"""Called at the start of each training batch."""
|
||||
pass
|
||||
|
||||
|
||||
def optimizer_step(trainer):
|
||||
"""Called when the optimizer takes a step."""
|
||||
pass
|
||||
|
||||
|
||||
def on_before_zero_grad(trainer):
|
||||
"""Called before the gradients are set to zero."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_batch_end(trainer):
|
||||
"""Called at the end of each training batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Called at the end of each training epoch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Called at the end of each fit epoch (train + val)."""
|
||||
pass
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Called when the model is saved."""
|
||||
pass
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Called when the training ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_params_update(trainer):
|
||||
"""Called when the model parameters are updated."""
|
||||
pass
|
||||
|
||||
|
||||
def teardown(trainer):
|
||||
"""Called during the teardown of the training process."""
|
||||
pass
|
||||
|
||||
|
||||
# Validator callbacks --------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_val_start(validator):
|
||||
"""Called when the validation starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_batch_start(validator):
|
||||
"""Called at the start of each validation batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_batch_end(validator):
|
||||
"""Called at the end of each validation batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_val_end(validator):
|
||||
"""Called when the validation ends."""
|
||||
pass
|
||||
|
||||
|
||||
# Predictor callbacks --------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_predict_start(predictor):
|
||||
"""Called when the prediction starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_batch_start(predictor):
|
||||
"""Called at the start of each prediction batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_batch_end(predictor):
|
||||
"""Called at the end of each prediction batch."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_postprocess_end(predictor):
|
||||
"""Called after the post-processing of the prediction ends."""
|
||||
pass
|
||||
|
||||
|
||||
def on_predict_end(predictor):
|
||||
"""Called when the prediction ends."""
|
||||
pass
|
||||
|
||||
|
||||
# Exporter callbacks ---------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def on_export_start(exporter):
|
||||
"""Called when the model export starts."""
|
||||
pass
|
||||
|
||||
|
||||
def on_export_end(exporter):
|
||||
"""Called when the model export ends."""
|
||||
pass
|
||||
|
||||
|
||||
@ -146,10 +179,23 @@ default_callbacks = {
|
||||
|
||||
|
||||
def get_default_callbacks():
|
||||
"""
|
||||
Return a copy of the default_callbacks dictionary with lists as default values.
|
||||
|
||||
Returns:
|
||||
(defaultdict): A defaultdict with keys from default_callbacks and empty lists as default values.
|
||||
"""
|
||||
return defaultdict(list, deepcopy(default_callbacks))
|
||||
|
||||
|
||||
def add_integration_callbacks(instance):
|
||||
"""
|
||||
Add integration callbacks from various sources to the instance's callbacks.
|
||||
|
||||
Args:
|
||||
instance (Trainer, Predictor, Validator, Exporter): An object with a 'callbacks' attribute that is a dictionary
|
||||
of callback lists.
|
||||
"""
|
||||
from .clearml import callbacks as clearml_callbacks
|
||||
from .comet import callbacks as comet_callbacks
|
||||
from .hub import callbacks as hub_callbacks
|
||||
|
||||
@ -59,6 +59,7 @@ def _log_plot(title, plot_path) -> None:
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Runs at start of pretraining routine; initializes and connects/ logs task to ClearML."""
|
||||
try:
|
||||
task = Task.current_task()
|
||||
if task:
|
||||
@ -83,11 +84,13 @@ def on_pretrain_routine_start(trainer):
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Logs debug samples for the first epoch of YOLO training."""
|
||||
if trainer.epoch == 1 and Task.current_task():
|
||||
_log_debug_samples(sorted(trainer.save_dir.glob('train_batch*.jpg')), 'Mosaic')
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Reports model information to logger at the end of an epoch."""
|
||||
task = Task.current_task()
|
||||
if task:
|
||||
# You should have access to the validation bboxes under jdict
|
||||
@ -105,12 +108,14 @@ def on_fit_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_val_end(validator):
|
||||
"""Logs validation results including labels and predictions."""
|
||||
if Task.current_task():
|
||||
# Log val_labels and val_pred
|
||||
_log_debug_samples(sorted(validator.save_dir.glob('val*.jpg')), 'Validation')
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Logs final model and its name on training completion."""
|
||||
task = Task.current_task()
|
||||
if task:
|
||||
# Log final results, CM matrix + PR plots
|
||||
|
||||
@ -36,6 +36,7 @@ _comet_image_prediction_count = 0
|
||||
|
||||
|
||||
def _get_experiment_type(mode, project_name):
|
||||
"""Return an experiment based on mode and project name."""
|
||||
if mode == 'offline':
|
||||
return comet_ml.OfflineExperiment(project_name=project_name)
|
||||
|
||||
@ -61,6 +62,7 @@ def _create_experiment(args):
|
||||
|
||||
|
||||
def _fetch_trainer_metadata(trainer):
|
||||
"""Returns metadata for YOLO training including epoch and asset saving status."""
|
||||
curr_epoch = trainer.epoch + 1
|
||||
|
||||
train_num_steps_per_epoch = len(trainer.train_loader.dataset) // trainer.batch_size
|
||||
@ -97,6 +99,7 @@ def _scale_bounding_box_to_original_image_shape(box, resized_image_shape, origin
|
||||
|
||||
|
||||
def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None):
|
||||
"""Format ground truth annotations for detection."""
|
||||
indices = batch['batch_idx'] == img_idx
|
||||
bboxes = batch['bboxes'][indices]
|
||||
if len(bboxes) == 0:
|
||||
@ -120,6 +123,7 @@ def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, c
|
||||
|
||||
|
||||
def _format_prediction_annotations_for_detection(image_path, metadata, class_label_map=None):
|
||||
"""Format YOLO predictions for object detection visualization."""
|
||||
stem = image_path.stem
|
||||
image_id = int(stem) if stem.isnumeric() else stem
|
||||
|
||||
@ -142,6 +146,7 @@ def _format_prediction_annotations_for_detection(image_path, metadata, class_lab
|
||||
|
||||
|
||||
def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, class_label_map):
|
||||
"""Join the ground truth and prediction annotations if they exist."""
|
||||
ground_truth_annotations = _format_ground_truth_annotations_for_detection(img_idx, image_path, batch,
|
||||
class_label_map)
|
||||
prediction_annotations = _format_prediction_annotations_for_detection(image_path, prediction_metadata_map,
|
||||
@ -153,6 +158,7 @@ def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, clas
|
||||
|
||||
|
||||
def _create_prediction_metadata_map(model_predictions):
|
||||
"""Create metadata map for model predictions by groupings them based on image ID."""
|
||||
pred_metadata_map = {}
|
||||
for prediction in model_predictions:
|
||||
pred_metadata_map.setdefault(prediction['image_id'], [])
|
||||
@ -162,6 +168,7 @@ def _create_prediction_metadata_map(model_predictions):
|
||||
|
||||
|
||||
def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch):
|
||||
"""Log the confusion matrix to Weights and Biases experiment."""
|
||||
conf_mat = trainer.validator.confusion_matrix.matrix
|
||||
names = list(trainer.data['names'].values()) + ['background']
|
||||
experiment.log_confusion_matrix(
|
||||
@ -174,6 +181,7 @@ def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch):
|
||||
|
||||
|
||||
def _log_images(experiment, image_paths, curr_step, annotations=None):
|
||||
"""Logs images to the experiment with optional annotations."""
|
||||
if annotations:
|
||||
for image_path, annotation in zip(image_paths, annotations):
|
||||
experiment.log_image(image_path, name=image_path.stem, step=curr_step, annotations=annotation)
|
||||
@ -184,6 +192,7 @@ def _log_images(experiment, image_paths, curr_step, annotations=None):
|
||||
|
||||
|
||||
def _log_image_predictions(experiment, validator, curr_step):
|
||||
"""Logs predicted boxes for a single image during training."""
|
||||
global _comet_image_prediction_count
|
||||
|
||||
task = validator.args.task
|
||||
@ -225,6 +234,7 @@ def _log_image_predictions(experiment, validator, curr_step):
|
||||
|
||||
|
||||
def _log_plots(experiment, trainer):
|
||||
"""Logs evaluation plots and label plots for the experiment."""
|
||||
plot_filenames = [trainer.save_dir / f'{plots}.png' for plots in EVALUATION_PLOT_NAMES]
|
||||
_log_images(experiment, plot_filenames, None)
|
||||
|
||||
@ -233,6 +243,7 @@ def _log_plots(experiment, trainer):
|
||||
|
||||
|
||||
def _log_model(experiment, trainer):
|
||||
"""Log the best-trained model to Comet.ml."""
|
||||
experiment.log_model(
|
||||
COMET_MODEL_NAME,
|
||||
file_or_folder=str(trainer.best),
|
||||
@ -242,12 +253,14 @@ def _log_model(experiment, trainer):
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Creates or resumes a CometML experiment at the start of a YOLO pre-training routine."""
|
||||
experiment = comet_ml.get_global_experiment()
|
||||
if not experiment:
|
||||
_create_experiment(trainer.args)
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Log metrics and save batch images at the end of training epochs."""
|
||||
experiment = comet_ml.get_global_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
@ -267,6 +280,7 @@ def on_train_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Logs model assets at the end of each epoch."""
|
||||
experiment = comet_ml.get_global_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
@ -296,6 +310,7 @@ def on_fit_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Perform operations at the end of training."""
|
||||
experiment = comet_ml.get_global_experiment()
|
||||
if not experiment:
|
||||
return
|
||||
|
||||
@ -9,6 +9,7 @@ from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Logs info before starting timer for upload rate limit."""
|
||||
session = getattr(trainer, 'hub_session', None)
|
||||
if session:
|
||||
# Start timer for upload rate limit
|
||||
@ -17,6 +18,7 @@ def on_pretrain_routine_end(trainer):
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Uploads training progress metrics at the end of each epoch."""
|
||||
session = getattr(trainer, 'hub_session', None)
|
||||
if session:
|
||||
# Upload metrics after val end
|
||||
@ -35,6 +37,7 @@ def on_fit_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Saves checkpoints to Ultralytics HUB with rate limiting."""
|
||||
session = getattr(trainer, 'hub_session', None)
|
||||
if session:
|
||||
# Upload checkpoints with rate limiting
|
||||
@ -46,6 +49,7 @@ def on_model_save(trainer):
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Upload final model and metrics to Ultralytics HUB at the end of training."""
|
||||
session = getattr(trainer, 'hub_session', None)
|
||||
if session:
|
||||
# Upload final model and metrics with exponential standoff
|
||||
@ -57,18 +61,22 @@ def on_train_end(trainer):
|
||||
|
||||
|
||||
def on_train_start(trainer):
|
||||
"""Run traces on train start."""
|
||||
traces(trainer.args, traces_sample_rate=1.0)
|
||||
|
||||
|
||||
def on_val_start(validator):
|
||||
"""Runs traces on validation start."""
|
||||
traces(validator.args, traces_sample_rate=1.0)
|
||||
|
||||
|
||||
def on_predict_start(predictor):
|
||||
"""Run traces on predict start."""
|
||||
traces(predictor.args, traces_sample_rate=1.0)
|
||||
|
||||
|
||||
def on_export_start(exporter):
|
||||
"""Run traces on export start."""
|
||||
traces(exporter.args, traces_sample_rate=1.0)
|
||||
|
||||
|
||||
|
||||
@ -16,6 +16,7 @@ except (ImportError, AssertionError):
|
||||
|
||||
|
||||
def on_pretrain_routine_end(trainer):
|
||||
"""Logs training parameters to MLflow."""
|
||||
global mlflow, run, run_id, experiment_name
|
||||
|
||||
if os.environ.get('MLFLOW_TRACKING_URI') is None:
|
||||
@ -45,17 +46,20 @@ def on_pretrain_routine_end(trainer):
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Logs training metrics to Mlflow."""
|
||||
if mlflow:
|
||||
metrics_dict = {f"{re.sub('[()]', '', k)}": float(v) for k, v in trainer.metrics.items()}
|
||||
run.log_metrics(metrics=metrics_dict, step=trainer.epoch)
|
||||
|
||||
|
||||
def on_model_save(trainer):
|
||||
"""Logs model and metrics to mlflow on save."""
|
||||
if mlflow:
|
||||
run.log_artifact(trainer.last)
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Called at end of train loop to log model artifact info."""
|
||||
if mlflow:
|
||||
root_dir = Path(__file__).resolve().parents[3]
|
||||
run.log_artifact(trainer.best)
|
||||
|
||||
@ -7,6 +7,7 @@ except (ImportError, AssertionError):
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Sends training metrics to Ray Tune at end of each epoch."""
|
||||
if ray.tune.is_session_enabled():
|
||||
metrics = trainer.metrics
|
||||
metrics['epoch'] = trainer.epoch
|
||||
|
||||
@ -12,12 +12,14 @@ writer = None # TensorBoard SummaryWriter instance
|
||||
|
||||
|
||||
def _log_scalars(scalars, step=0):
|
||||
"""Logs scalar values to TensorBoard."""
|
||||
if writer:
|
||||
for k, v in scalars.items():
|
||||
writer.add_scalar(k, v, step)
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Initialize TensorBoard logging with SummaryWriter."""
|
||||
if SummaryWriter:
|
||||
try:
|
||||
global writer
|
||||
@ -29,10 +31,12 @@ def on_pretrain_routine_start(trainer):
|
||||
|
||||
|
||||
def on_batch_end(trainer):
|
||||
"""Logs scalar statistics at the end of a training batch."""
|
||||
_log_scalars(trainer.label_loss_items(trainer.tloss, prefix='train'), trainer.epoch + 1)
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Logs epoch metrics at end of training epoch."""
|
||||
_log_scalars(trainer.metrics, trainer.epoch + 1)
|
||||
|
||||
|
||||
|
||||
@ -11,11 +11,13 @@ except (ImportError, AssertionError):
|
||||
|
||||
|
||||
def on_pretrain_routine_start(trainer):
|
||||
"""Initiate and start project if module is present."""
|
||||
wb.init(project=trainer.args.project or 'YOLOv8', name=trainer.args.name, config=vars(
|
||||
trainer.args)) if not wb.run else wb.run
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""Logs training metrics and model information at the end of an epoch."""
|
||||
wb.run.log(trainer.metrics, step=trainer.epoch + 1)
|
||||
if trainer.epoch == 0:
|
||||
model_info = {
|
||||
@ -26,6 +28,7 @@ def on_fit_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_train_epoch_end(trainer):
|
||||
"""Log metrics and save images at the end of each training epoch."""
|
||||
wb.run.log(trainer.label_loss_items(trainer.tloss, prefix='train'), step=trainer.epoch + 1)
|
||||
wb.run.log(trainer.lr, step=trainer.epoch + 1)
|
||||
if trainer.epoch == 1:
|
||||
@ -35,6 +38,7 @@ def on_train_epoch_end(trainer):
|
||||
|
||||
|
||||
def on_train_end(trainer):
|
||||
"""Save the best model as an artifact at end of training."""
|
||||
art = wb.Artifact(type='model', name=f'run_{wb.run.id}_model')
|
||||
if trainer.best.exists():
|
||||
art.add_file(trainer.best)
|
||||
|
||||
@ -295,7 +295,7 @@ def check_file(file, suffix='', download=True, hard=True):
|
||||
|
||||
|
||||
def check_yaml(file, suffix=('.yaml', '.yml'), hard=True):
|
||||
# Search/download YAML file (if necessary) and return path, checking suffix
|
||||
"""Search/download YAML file (if necessary) and return path, checking suffix."""
|
||||
return check_file(file, suffix, hard=hard)
|
||||
|
||||
|
||||
@ -315,6 +315,7 @@ def check_imshow(warn=False):
|
||||
|
||||
|
||||
def check_yolo(verbose=True, device=''):
|
||||
"""Return a human-readable YOLO software and hardware summary."""
|
||||
from ultralytics.yolo.utils.torch_utils import select_device
|
||||
|
||||
if is_colab():
|
||||
|
||||
@ -24,6 +24,7 @@ def find_free_network_port() -> int:
|
||||
|
||||
|
||||
def generate_ddp_file(trainer):
|
||||
"""Generates a DDP file and returns its file name."""
|
||||
module, name = f'{trainer.__class__.__module__}.{trainer.__class__.__name__}'.rsplit('.', 1)
|
||||
|
||||
content = f'''cfg = {vars(trainer.args)} \nif __name__ == "__main__":
|
||||
@ -43,6 +44,7 @@ def generate_ddp_file(trainer):
|
||||
|
||||
|
||||
def generate_ddp_command(world_size, trainer):
|
||||
"""Generates and returns command for distributed training."""
|
||||
import __main__ # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
|
||||
if not trainer.resume:
|
||||
shutil.rmtree(trainer.save_dir) # remove the save_dir
|
||||
|
||||
@ -192,7 +192,7 @@ def attempt_download_asset(file, repo='ultralytics/assets', release='v0.0.0'):
|
||||
|
||||
|
||||
def download(url, dir=Path.cwd(), unzip=True, delete=False, curl=False, threads=1, retry=3):
|
||||
# Multithreaded file download and unzip function, used in data.yaml for autodownload
|
||||
"""Downloads and unzips files concurrently if threads > 1, else sequentially."""
|
||||
dir = Path(dir)
|
||||
dir.mkdir(parents=True, exist_ok=True) # make directory
|
||||
if threads > 1:
|
||||
|
||||
@ -6,4 +6,5 @@ from ultralytics.yolo.utils import emojis
|
||||
class HUBModelError(Exception):
|
||||
|
||||
def __init__(self, message='Model not found. Please check model URL and try again.'):
|
||||
"""Create an exception for when a model is not found."""
|
||||
super().__init__(emojis(message))
|
||||
|
||||
@ -11,13 +11,16 @@ class WorkingDirectory(contextlib.ContextDecorator):
|
||||
"""Usage: @WorkingDirectory(dir) decorator or 'with WorkingDirectory(dir):' context manager."""
|
||||
|
||||
def __init__(self, new_dir):
|
||||
"""Sets the working directory to 'new_dir' upon instantiation."""
|
||||
self.dir = new_dir # new dir
|
||||
self.cwd = Path.cwd().resolve() # current dir
|
||||
|
||||
def __enter__(self):
|
||||
"""Changes the current directory to the specified directory."""
|
||||
os.chdir(self.dir)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Restore the current working directory on context exit."""
|
||||
os.chdir(self.cwd)
|
||||
|
||||
|
||||
|
||||
@ -14,6 +14,7 @@ def _ntuple(n):
|
||||
"""From PyTorch internals."""
|
||||
|
||||
def parse(x):
|
||||
"""Parse bounding boxes format between XYWH and LTWH."""
|
||||
return x if isinstance(x, abc.Iterable) else tuple(repeat(x, n))
|
||||
|
||||
return parse
|
||||
@ -64,6 +65,7 @@ class Bboxes:
|
||||
# return Bboxes(bboxes, format)
|
||||
|
||||
def convert(self, format):
|
||||
"""Converts bounding box format from one type to another."""
|
||||
assert format in _formats
|
||||
if self.format == format:
|
||||
return
|
||||
@ -77,6 +79,7 @@ class Bboxes:
|
||||
self.format = format
|
||||
|
||||
def areas(self):
|
||||
"""Return box areas."""
|
||||
self.convert('xyxy')
|
||||
return (self.bboxes[:, 2] - self.bboxes[:, 0]) * (self.bboxes[:, 3] - self.bboxes[:, 1])
|
||||
|
||||
@ -125,6 +128,7 @@ class Bboxes:
|
||||
self.bboxes[:, 3] += offset[3]
|
||||
|
||||
def __len__(self):
|
||||
"""Return the number of boxes."""
|
||||
return len(self.bboxes)
|
||||
|
||||
@classmethod
|
||||
@ -202,9 +206,11 @@ class Instances:
|
||||
self.segments = segments
|
||||
|
||||
def convert_bbox(self, format):
|
||||
"""Convert bounding box format."""
|
||||
self._bboxes.convert(format=format)
|
||||
|
||||
def bbox_areas(self):
|
||||
"""Calculate the area of bounding boxes."""
|
||||
self._bboxes.areas()
|
||||
|
||||
def scale(self, scale_w, scale_h, bbox_only=False):
|
||||
@ -219,6 +225,7 @@ class Instances:
|
||||
self.keypoints[..., 1] *= scale_h
|
||||
|
||||
def denormalize(self, w, h):
|
||||
"""Denormalizes boxes, segments, and keypoints from normalized coordinates."""
|
||||
if not self.normalized:
|
||||
return
|
||||
self._bboxes.mul(scale=(w, h, w, h))
|
||||
@ -230,6 +237,7 @@ class Instances:
|
||||
self.normalized = False
|
||||
|
||||
def normalize(self, w, h):
|
||||
"""Normalize bounding boxes, segments, and keypoints to image dimensions."""
|
||||
if self.normalized:
|
||||
return
|
||||
self._bboxes.mul(scale=(1 / w, 1 / h, 1 / w, 1 / h))
|
||||
@ -279,6 +287,7 @@ class Instances:
|
||||
)
|
||||
|
||||
def flipud(self, h):
|
||||
"""Flips the coordinates of bounding boxes, segments, and keypoints vertically."""
|
||||
if self._bboxes.format == 'xyxy':
|
||||
y1 = self.bboxes[:, 1].copy()
|
||||
y2 = self.bboxes[:, 3].copy()
|
||||
@ -291,6 +300,7 @@ class Instances:
|
||||
self.keypoints[..., 1] = h - self.keypoints[..., 1]
|
||||
|
||||
def fliplr(self, w):
|
||||
"""Reverses the order of the bounding boxes and segments horizontally."""
|
||||
if self._bboxes.format == 'xyxy':
|
||||
x1 = self.bboxes[:, 0].copy()
|
||||
x2 = self.bboxes[:, 2].copy()
|
||||
@ -303,6 +313,7 @@ class Instances:
|
||||
self.keypoints[..., 0] = w - self.keypoints[..., 0]
|
||||
|
||||
def clip(self, w, h):
|
||||
"""Clips bounding boxes, segments, and keypoints values to stay within image boundaries."""
|
||||
ori_format = self._bboxes.format
|
||||
self.convert_bbox(format='xyxy')
|
||||
self.bboxes[:, [0, 2]] = self.bboxes[:, [0, 2]].clip(0, w)
|
||||
@ -316,6 +327,7 @@ class Instances:
|
||||
self.keypoints[..., 1] = self.keypoints[..., 1].clip(0, h)
|
||||
|
||||
def update(self, bboxes, segments=None, keypoints=None):
|
||||
"""Updates instance variables."""
|
||||
new_bboxes = Bboxes(bboxes, format=self._bboxes.format)
|
||||
self._bboxes = new_bboxes
|
||||
if segments is not None:
|
||||
@ -324,6 +336,7 @@ class Instances:
|
||||
self.keypoints = keypoints
|
||||
|
||||
def __len__(self):
|
||||
"""Return the length of the instance list."""
|
||||
return len(self.bboxes)
|
||||
|
||||
@classmethod
|
||||
@ -363,4 +376,5 @@ class Instances:
|
||||
|
||||
@property
|
||||
def bboxes(self):
|
||||
"""Return bounding boxes."""
|
||||
return self._bboxes.bboxes
|
||||
|
||||
@ -12,9 +12,11 @@ class VarifocalLoss(nn.Module):
|
||||
"""Varifocal loss by Zhang et al. https://arxiv.org/abs/2008.13367."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the VarifocalLoss class."""
|
||||
super().__init__()
|
||||
|
||||
def forward(self, pred_score, gt_score, label, alpha=0.75, gamma=2.0):
|
||||
"""Computes varfocal loss."""
|
||||
weight = alpha * pred_score.sigmoid().pow(gamma) * (1 - label) + gt_score * label
|
||||
with torch.cuda.amp.autocast(enabled=False):
|
||||
loss = (F.binary_cross_entropy_with_logits(pred_score.float(), gt_score.float(), reduction='none') *
|
||||
@ -25,6 +27,7 @@ class VarifocalLoss(nn.Module):
|
||||
class BboxLoss(nn.Module):
|
||||
|
||||
def __init__(self, reg_max, use_dfl=False):
|
||||
"""Initialize the BboxLoss module with regularization maximum and DFL settings."""
|
||||
super().__init__()
|
||||
self.reg_max = reg_max
|
||||
self.use_dfl = use_dfl
|
||||
@ -64,6 +67,7 @@ class KeypointLoss(nn.Module):
|
||||
self.sigmas = sigmas
|
||||
|
||||
def forward(self, pred_kpts, gt_kpts, kpt_mask, area):
|
||||
"""Calculates keypoint loss factor and Euclidean distance loss for predicted and actual keypoints."""
|
||||
d = (pred_kpts[..., 0] - gt_kpts[..., 0]) ** 2 + (pred_kpts[..., 1] - gt_kpts[..., 1]) ** 2
|
||||
kpt_loss_factor = (torch.sum(kpt_mask != 0) + torch.sum(kpt_mask == 0)) / (torch.sum(kpt_mask != 0) + 1e-9)
|
||||
# e = d / (2 * (area * self.sigmas) ** 2 + 1e-9) # from formula
|
||||
|
||||
@ -180,6 +180,7 @@ class FocalLoss(nn.Module):
|
||||
"""Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)."""
|
||||
|
||||
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
|
||||
"""Initialize FocalLoss object with given loss function and hyperparameters."""
|
||||
super().__init__()
|
||||
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
|
||||
self.gamma = gamma
|
||||
@ -188,6 +189,7 @@ class FocalLoss(nn.Module):
|
||||
self.loss_fcn.reduction = 'none' # required to apply FL to each element
|
||||
|
||||
def forward(self, pred, true):
|
||||
"""Calculates and updates confusion matrix for object detection/classification tasks."""
|
||||
loss = self.loss_fcn(pred, true)
|
||||
# p_t = torch.exp(-loss)
|
||||
# loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
|
||||
@ -220,6 +222,7 @@ class ConfusionMatrix:
|
||||
"""
|
||||
|
||||
def __init__(self, nc, conf=0.25, iou_thres=0.45, task='detect'):
|
||||
"""Initialize attributes for the YOLO model."""
|
||||
self.task = task
|
||||
self.matrix = np.zeros((nc + 1, nc + 1)) if self.task == 'detect' else np.zeros((nc, nc))
|
||||
self.nc = nc # number of classes
|
||||
@ -285,9 +288,11 @@ class ConfusionMatrix:
|
||||
self.matrix[dc, self.nc] += 1 # predicted background
|
||||
|
||||
def matrix(self):
|
||||
"""Returns the confusion matrix."""
|
||||
return self.matrix
|
||||
|
||||
def tp_fp(self):
|
||||
"""Returns true positives and false positives."""
|
||||
tp = self.matrix.diagonal() # true positives
|
||||
fp = self.matrix.sum(1) - tp # false positives
|
||||
# fn = self.matrix.sum(0) - tp # false negatives (missed detections)
|
||||
@ -679,6 +684,7 @@ class DetMetrics(SimpleClass):
|
||||
self.speed = {'preprocess': 0.0, 'inference': 0.0, 'loss': 0.0, 'postprocess': 0.0}
|
||||
|
||||
def process(self, tp, conf, pred_cls, target_cls):
|
||||
"""Process predicted results for object detection and update metrics."""
|
||||
results = ap_per_class(tp, conf, pred_cls, target_cls, plot=self.plot, save_dir=self.save_dir,
|
||||
names=self.names)[2:]
|
||||
self.box.nc = len(self.names)
|
||||
@ -686,28 +692,35 @@ class DetMetrics(SimpleClass):
|
||||
|
||||
@property
|
||||
def keys(self):
|
||||
"""Returns a list of keys for accessing specific metrics."""
|
||||
return ['metrics/precision(B)', 'metrics/recall(B)', 'metrics/mAP50(B)', 'metrics/mAP50-95(B)']
|
||||
|
||||
def mean_results(self):
|
||||
"""Calculate mean of detected objects & return precision, recall, mAP50, and mAP50-95."""
|
||||
return self.box.mean_results()
|
||||
|
||||
def class_result(self, i):
|
||||
"""Return the result of evaluating the performance of an object detection model on a specific class."""
|
||||
return self.box.class_result(i)
|
||||
|
||||
@property
|
||||
def maps(self):
|
||||
"""Returns mean Average Precision (mAP) scores per class."""
|
||||
return self.box.maps
|
||||
|
||||
@property
|
||||
def fitness(self):
|
||||
"""Returns the fitness of box object."""
|
||||
return self.box.fitness()
|
||||
|
||||
@property
|
||||
def ap_class_index(self):
|
||||
"""Returns the average precision index per class."""
|
||||
return self.box.ap_class_index
|
||||
|
||||
@property
|
||||
def results_dict(self):
|
||||
"""Returns dictionary of computed performance metrics and statistics."""
|
||||
return dict(zip(self.keys + ['fitness'], self.mean_results() + [self.fitness]))
|
||||
|
||||
|
||||
@ -781,22 +794,27 @@ class SegmentMetrics(SimpleClass):
|
||||
|
||||
@property
|
||||
def keys(self):
|
||||
"""Returns a list of keys for accessing metrics."""
|
||||
return [
|
||||
'metrics/precision(B)', 'metrics/recall(B)', 'metrics/mAP50(B)', 'metrics/mAP50-95(B)',
|
||||
'metrics/precision(M)', 'metrics/recall(M)', 'metrics/mAP50(M)', 'metrics/mAP50-95(M)']
|
||||
|
||||
def mean_results(self):
|
||||
"""Return the mean metrics for bounding box and segmentation results."""
|
||||
return self.box.mean_results() + self.seg.mean_results()
|
||||
|
||||
def class_result(self, i):
|
||||
"""Returns classification results for a specified class index."""
|
||||
return self.box.class_result(i) + self.seg.class_result(i)
|
||||
|
||||
@property
|
||||
def maps(self):
|
||||
"""Returns mAP scores for object detection and semantic segmentation models."""
|
||||
return self.box.maps + self.seg.maps
|
||||
|
||||
@property
|
||||
def fitness(self):
|
||||
"""Get the fitness score for both segmentation and bounding box models."""
|
||||
return self.seg.fitness() + self.box.fitness()
|
||||
|
||||
@property
|
||||
@ -806,6 +824,7 @@ class SegmentMetrics(SimpleClass):
|
||||
|
||||
@property
|
||||
def results_dict(self):
|
||||
"""Returns results of object detection model for evaluation."""
|
||||
return dict(zip(self.keys + ['fitness'], self.mean_results() + [self.fitness]))
|
||||
|
||||
|
||||
@ -846,6 +865,7 @@ class PoseMetrics(SegmentMetrics):
|
||||
self.speed = {'preprocess': 0.0, 'inference': 0.0, 'loss': 0.0, 'postprocess': 0.0}
|
||||
|
||||
def __getattr__(self, attr):
|
||||
"""Raises an AttributeError if an invalid attribute is accessed."""
|
||||
name = self.__class__.__name__
|
||||
raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
|
||||
|
||||
@ -884,22 +904,27 @@ class PoseMetrics(SegmentMetrics):
|
||||
|
||||
@property
|
||||
def keys(self):
|
||||
"""Returns list of evaluation metric keys."""
|
||||
return [
|
||||
'metrics/precision(B)', 'metrics/recall(B)', 'metrics/mAP50(B)', 'metrics/mAP50-95(B)',
|
||||
'metrics/precision(P)', 'metrics/recall(P)', 'metrics/mAP50(P)', 'metrics/mAP50-95(P)']
|
||||
|
||||
def mean_results(self):
|
||||
"""Return the mean results of box and pose."""
|
||||
return self.box.mean_results() + self.pose.mean_results()
|
||||
|
||||
def class_result(self, i):
|
||||
"""Return the class-wise detection results for a specific class i."""
|
||||
return self.box.class_result(i) + self.pose.class_result(i)
|
||||
|
||||
@property
|
||||
def maps(self):
|
||||
"""Returns the mean average precision (mAP) per class for both box and pose detections."""
|
||||
return self.box.maps + self.pose.maps
|
||||
|
||||
@property
|
||||
def fitness(self):
|
||||
"""Computes classification metrics and speed using the `targets` and `pred` inputs."""
|
||||
return self.pose.fitness() + self.box.fitness()
|
||||
|
||||
|
||||
@ -935,12 +960,15 @@ class ClassifyMetrics(SimpleClass):
|
||||
|
||||
@property
|
||||
def fitness(self):
|
||||
"""Returns top-5 accuracy as fitness score."""
|
||||
return self.top5
|
||||
|
||||
@property
|
||||
def results_dict(self):
|
||||
"""Returns a dictionary with model's performance metrics and fitness score."""
|
||||
return dict(zip(self.keys + ['fitness'], [self.top1, self.top5, self.fitness]))
|
||||
|
||||
@property
|
||||
def keys(self):
|
||||
"""Returns a list of keys for the results_dict property."""
|
||||
return ['metrics/accuracy_top1', 'metrics/accuracy_top5']
|
||||
|
||||
@ -33,6 +33,7 @@ class Colors:
|
||||
dtype=np.uint8)
|
||||
|
||||
def __call__(self, i, bgr=False):
|
||||
"""Converts hex color codes to rgb values."""
|
||||
c = self.palette[int(i) % self.n]
|
||||
return (c[2], c[1], c[0]) if bgr else c
|
||||
|
||||
@ -47,6 +48,7 @@ colors = Colors() # create instance for 'from utils.plots import colors'
|
||||
class Annotator:
|
||||
# YOLOv8 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
||||
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
||||
"""Initialize the Annotator class with image and line width along with color palette for keypoints and limbs."""
|
||||
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
||||
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
||||
self.pil = pil or non_ascii
|
||||
@ -71,7 +73,7 @@ class Annotator:
|
||||
self.kpt_color = colors.pose_palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]
|
||||
|
||||
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
||||
# Add one xyxy box to image with label
|
||||
"""Add one xyxy box to image with label."""
|
||||
if isinstance(box, torch.Tensor):
|
||||
box = box.tolist()
|
||||
if self.pil or not is_ascii(label):
|
||||
@ -191,7 +193,7 @@ class Annotator:
|
||||
self.draw.rectangle(xy, fill, outline, width)
|
||||
|
||||
def text(self, xy, text, txt_color=(255, 255, 255), anchor='top'):
|
||||
# Add text to image (PIL-only)
|
||||
"""Adds text to an image using PIL or cv2."""
|
||||
if anchor == 'bottom': # start y from font bottom
|
||||
w, h = self.font.getsize(text) # text width, height
|
||||
xy[1] += 1 - h
|
||||
@ -214,6 +216,7 @@ class Annotator:
|
||||
@TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395
|
||||
@plt_settings()
|
||||
def plot_labels(boxes, cls, names=(), save_dir=Path('')):
|
||||
"""Save and plot image with no axis or spines."""
|
||||
import pandas as pd
|
||||
import seaborn as sn
|
||||
|
||||
@ -260,7 +263,7 @@ def plot_labels(boxes, cls, names=(), save_dir=Path('')):
|
||||
|
||||
|
||||
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
|
||||
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
||||
"""Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop."""
|
||||
b = xyxy2xywh(xyxy.view(-1, 4)) # boxes
|
||||
if square:
|
||||
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
||||
|
||||
@ -69,6 +69,7 @@ class TaskAlignedAssigner(nn.Module):
|
||||
"""
|
||||
|
||||
def __init__(self, topk=13, num_classes=80, alpha=1.0, beta=6.0, eps=1e-9):
|
||||
"""Initialize a TaskAlignedAssigner object with customizable hyperparameters."""
|
||||
super().__init__()
|
||||
self.topk = topk
|
||||
self.num_classes = num_classes
|
||||
@ -137,6 +138,7 @@ class TaskAlignedAssigner(nn.Module):
|
||||
return mask_pos, align_metric, overlaps
|
||||
|
||||
def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt):
|
||||
"""Compute alignment metric given predicted and ground truth bounding boxes."""
|
||||
na = pd_bboxes.shape[-2]
|
||||
mask_gt = mask_gt.bool() # b, max_num_obj, h*w
|
||||
overlaps = torch.zeros([self.bs, self.n_max_boxes, na], dtype=pd_bboxes.dtype, device=pd_bboxes.device)
|
||||
|
||||
@ -43,6 +43,7 @@ def smart_inference_mode():
|
||||
"""Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator."""
|
||||
|
||||
def decorate(fn):
|
||||
"""Applies appropriate torch decorator for inference mode based on torch version."""
|
||||
return (torch.inference_mode if TORCH_1_9 else torch.no_grad)()(fn)
|
||||
|
||||
return decorate
|
||||
@ -232,7 +233,7 @@ def make_divisible(x, divisor):
|
||||
|
||||
|
||||
def copy_attr(a, b, include=(), exclude=()):
|
||||
# Copy attributes from 'b' to 'a', options to only include [...] and to exclude [...]
|
||||
"""Copies attributes from object 'b' to object 'a', with options to include/exclude certain attributes."""
|
||||
for k, v in b.__dict__.items():
|
||||
if (len(include) and k not in include) or k.startswith('_') or k in exclude:
|
||||
continue
|
||||
@ -246,7 +247,7 @@ def get_latest_opset():
|
||||
|
||||
|
||||
def intersect_dicts(da, db, exclude=()):
|
||||
# Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
|
||||
"""Returns a dictionary of intersecting keys with matching shapes, excluding 'exclude' keys, using da values."""
|
||||
return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
|
||||
|
||||
|
||||
@ -310,7 +311,7 @@ class ModelEMA:
|
||||
# assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype}, model {msd[k].dtype}'
|
||||
|
||||
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
||||
# Update EMA attributes
|
||||
"""Updates attributes and saves stripped model with optimizer removed."""
|
||||
if self.enabled:
|
||||
copy_attr(self.ema, model, include, exclude)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user