Add docstrings to YOLOv5 functions (#12760)

* Add docstrings to top level files * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Add docstrings * Auto-format by https://ultralytics.com/actions * Update activations.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Auto-format by https://ultralytics.com/actions --------- Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
2024-02-25 14:04:01 +01:00
1 parent ca00c34c2a
commit 41603da162
@@ -73,6 +73,7 @@ def _json_default(value):
 class Loggers:
    # YOLOv5 Loggers class
    def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
+        """Initializes loggers for YOLOv5 training and validation metrics, paths, and options."""
        self.save_dir = save_dir
        self.weights = weights
        self.opt = opt
@@ -150,7 +151,7 @@ class Loggers:

    @property
    def remote_dataset(self):
-        # Get data_dict if custom dataset artifact link is provided
+        """Fetches dataset dictionary from remote logging services like ClearML, Weights & Biases, or Comet ML."""
        data_dict = None
        if self.clearml:
            data_dict = self.clearml.data_dict
@@ -162,15 +163,17 @@ class Loggers:
        return data_dict

    def on_train_start(self):
+        """Initializes the training process for Comet ML logger if it's configured."""
        if self.comet_logger:
            self.comet_logger.on_train_start()

    def on_pretrain_routine_start(self):
+        """Invokes pre-training routine start hook for Comet ML logger if available."""
        if self.comet_logger:
            self.comet_logger.on_pretrain_routine_start()

    def on_pretrain_routine_end(self, labels, names):
-        # Callback runs on pre-train routine end
+        """Callback that runs at the end of pre-training routine, logging label plots if enabled."""
        if self.plots:
            plot_labels(labels, names, self.save_dir)
            paths = self.save_dir.glob("*labels*.jpg")  # training labels
@@ -183,6 +186,7 @@ class Loggers:
                    self.clearml.log_plot(title=path.stem, plot_path=path)

    def on_train_batch_end(self, model, ni, imgs, targets, paths, vals):
+        """Logs training batch end events, plots images, and updates external loggers with batch-end data."""
        log_dict = dict(zip(self.keys[:3], vals))
        # Callback runs on train batch end
        # ni: number integrated batches (since train start)
@@ -203,7 +207,7 @@ class Loggers:
            self.comet_logger.on_train_batch_end(log_dict, step=ni)

    def on_train_epoch_end(self, epoch):
-        # Callback runs on train epoch end
+        """Callback that updates the current epoch in Weights & Biases at the end of a training epoch."""
        if self.wandb:
            self.wandb.current_epoch = epoch + 1

@@ -211,22 +215,24 @@ class Loggers:
            self.comet_logger.on_train_epoch_end(epoch)

    def on_val_start(self):
+        """Callback that signals the start of a validation phase to the Comet logger."""
        if self.comet_logger:
            self.comet_logger.on_val_start()

    def on_val_image_end(self, pred, predn, path, names, im):
-        # Callback runs on val image end
+        """Callback that logs a validation image and its predictions to WandB or ClearML."""
        if self.wandb:
            self.wandb.val_one_image(pred, predn, path, names, im)
        if self.clearml:
            self.clearml.log_image_with_boxes(path, pred, names, im)

    def on_val_batch_end(self, batch_i, im, targets, paths, shapes, out):
+        """Logs validation batch results to Comet ML during training at the end of each validation batch."""
        if self.comet_logger:
            self.comet_logger.on_val_batch_end(batch_i, im, targets, paths, shapes, out)

    def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
-        # Callback runs on val end
+        """Logs validation results to WandB or ClearML at the end of the validation process."""
        if self.wandb or self.clearml:
            files = sorted(self.save_dir.glob("val*.jpg"))
        if self.wandb:
@@ -238,7 +244,7 @@ class Loggers:
            self.comet_logger.on_val_end(nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)

    def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
-        # Callback runs at the end of each fit (train+val) epoch
+        """Callback that logs metrics and saves them to CSV or NDJSON at the end of each fit (train+val) epoch."""
        x = dict(zip(self.keys, vals))
        if self.csv:
            file = self.save_dir / "results.csv"
@@ -277,7 +283,7 @@ class Loggers:
            self.comet_logger.on_fit_epoch_end(x, epoch=epoch)

    def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
-        # Callback runs on model save event
+        """Callback that handles model saving events, logging to Weights & Biases or ClearML if enabled."""
        if (epoch + 1) % self.opt.save_period == 0 and not final_epoch and self.opt.save_period != -1:
            if self.wandb:
                self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
@@ -290,7 +296,7 @@ class Loggers:
            self.comet_logger.on_model_save(last, epoch, final_epoch, best_fitness, fi)

    def on_train_end(self, last, best, epoch, results):
-        # Callback runs on training end, i.e. saving best model
+        """Callback that runs at the end of training to save plots and log results."""
        if self.plots:
            plot_results(file=self.save_dir / "results.csv")  # save results.png
        files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))]
@@ -326,7 +332,7 @@ class Loggers:
            self.comet_logger.on_train_end(files, self.save_dir, last, best, epoch, final_results)

    def on_params_update(self, params: dict):
-        # Update hyperparams or configs of the experiment
+        """Updates experiment hyperparameters or configurations in WandB, Comet, or ClearML."""
        if self.wandb:
            self.wandb.wandb_run.config.update(params, allow_val_change=True)
        if self.comet_logger:
@@ -346,7 +352,7 @@ class GenericLogger:
    """

    def __init__(self, opt, console_logger, include=("tb", "wandb", "clearml")):
-        # init default loggers
+        """Initializes a generic logger with optional TensorBoard, W&B, and ClearML support."""
        self.save_dir = Path(opt.save_dir)
        self.include = include
        self.console_logger = console_logger
@@ -381,7 +387,7 @@ class GenericLogger:
            self.clearml = None

    def log_metrics(self, metrics, epoch):
-        # Log metrics dictionary to all loggers
+        """Logs metrics to CSV, TensorBoard, W&B, and ClearML; `metrics` is a dict, `epoch` is an int."""
        if self.csv:
            keys, vals = list(metrics.keys()), list(metrics.values())
            n = len(metrics) + 1  # number of cols
@@ -400,7 +406,7 @@ class GenericLogger:
            self.clearml.log_scalars(metrics, epoch)

    def log_images(self, files, name="Images", epoch=0):
-        # Log images to all loggers
+        """Logs images to all loggers with optional naming and epoch specification."""
        files = [Path(f) for f in (files if isinstance(files, (tuple, list)) else [files])]  # to Path
        files = [f for f in files if f.exists()]  # filter by exists

@@ -418,11 +424,12 @@ class GenericLogger:
                self.clearml.log_debug_samples(files, title=name)

    def log_graph(self, model, imgsz=(640, 640)):
-        # Log model graph to all loggers
+        """Logs model graph to all configured loggers with specified input image size."""
        if self.tb:
            log_tensorboard_graph(self.tb, model, imgsz)

    def log_model(self, model_path, epoch=0, metadata=None):
+        """Logs the model to all configured loggers with optional epoch and metadata."""
        if metadata is None:
            metadata = {}
        # Log model to all loggers
@@ -434,7 +441,7 @@ class GenericLogger:
            self.clearml.log_model(model_path=model_path, model_name=model_path.stem)

    def update_params(self, params):
-        # Update the parameters logged
+        """Updates logged parameters in WandB and/or ClearML if enabled."""
        if self.wandb:
            wandb.run.config.update(params, allow_val_change=True)
        if self.clearml:
@@ -442,7 +449,7 @@ class GenericLogger:


 def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
-    # Log model graph to TensorBoard
+    """Logs the model graph to TensorBoard with specified image size and model."""
    try:
        p = next(model.parameters())  # for device, type
        imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz  # expand
@@ -455,7 +462,7 @@ def log_tensorboard_graph(tb, model, imgsz=(640, 640)):


 def web_project_name(project):
-    # Convert local project name to web project name
+    """Converts a local project name to a standardized web project name with optional suffixes."""
    if not project.startswith("runs/train"):
        return project
    suffix = "-Classify" if project.endswith("-cls") else "-Segment" if project.endswith("-seg") else ""
@@ -165,6 +165,7 @@ class CometLogger:
            self.experiment.log_other("optimizer_parameters", json.dumps(self.hyp))

    def _get_experiment(self, mode, experiment_id=None):
+        """Returns a new or existing Comet.ml experiment based on mode and optional experiment_id."""
        if mode == "offline":
            return (
                comet_ml.ExistingOfflineExperiment(
@@ -197,21 +198,27 @@ class CometLogger:
        return

    def log_metrics(self, log_dict, **kwargs):
+        """Logs metrics to the current experiment, accepting a dictionary of metric names and values."""
        self.experiment.log_metrics(log_dict, **kwargs)

    def log_parameters(self, log_dict, **kwargs):
+        """Logs parameters to the current experiment, accepting a dictionary of parameter names and values."""
        self.experiment.log_parameters(log_dict, **kwargs)

    def log_asset(self, asset_path, **kwargs):
+        """Logs a file or directory as an asset to the current experiment."""
        self.experiment.log_asset(asset_path, **kwargs)

    def log_asset_data(self, asset, **kwargs):
+        """Logs in-memory data as an asset to the current experiment, with optional kwargs."""
        self.experiment.log_asset_data(asset, **kwargs)

    def log_image(self, img, **kwargs):
+        """Logs an image to the current experiment with optional kwargs."""
        self.experiment.log_image(img, **kwargs)

    def log_model(self, path, opt, epoch, fitness_score, best_model=False):
+        """Logs model checkpoint to experiment with path, options, epoch, fitness, and best model flag."""
        if not self.save_model:
            return

@@ -235,6 +242,7 @@ class CometLogger:
            )

    def check_dataset(self, data_file):
+        """Validates the dataset configuration by loading the YAML file specified in `data_file`."""
        with open(data_file) as f:
            data_config = yaml.safe_load(f)

@@ -247,6 +255,7 @@ class CometLogger:
        return check_dataset(data_file)

    def log_predictions(self, image, labelsn, path, shape, predn):
+        """Logs predictions with IOU filtering, given image, labels, path, shape, and predictions."""
        if self.logged_images_count >= self.max_images:
            return
        detections = predn[predn[:, 4] > self.conf_thres]
@@ -287,6 +296,7 @@ class CometLogger:
        return

    def preprocess_prediction(self, image, labels, shape, pred):
+        """Processes prediction data, resizing labels and adding dataset metadata."""
        nl, _ = labels.shape[0], pred.shape[0]

        # Predictions
@@ -306,6 +316,7 @@ class CometLogger:
        return predn, labelsn

    def add_assets_to_artifact(self, artifact, path, asset_path, split):
+        """Adds image and label assets to a wandb artifact given dataset split and paths."""
        img_paths = sorted(glob.glob(f"{asset_path}/*"))
        label_paths = img2label_paths(img_paths)

@@ -331,6 +342,7 @@ class CometLogger:
        return artifact

    def upload_dataset_artifact(self):
+        """Uploads a YOLOv5 dataset as an artifact to the Comet.ml platform."""
        dataset_name = self.data_dict.get("dataset_name", "yolov5-dataset")
        path = str((ROOT / Path(self.data_dict["path"])).resolve())

@@ -355,6 +367,7 @@ class CometLogger:
        return

    def download_dataset_artifact(self, artifact_path):
+        """Downloads a dataset artifact to a specified directory using the experiment's logged artifact."""
        logged_artifact = self.experiment.get_artifact(artifact_path)
        artifact_save_dir = str(Path(self.opt.save_dir) / logged_artifact.name)
        logged_artifact.download(artifact_save_dir)
@@ -374,6 +387,7 @@ class CometLogger:
        return self.update_data_paths(data_dict)

    def update_data_paths(self, data_dict):
+        """Updates data paths in the dataset dictionary, defaulting 'path' to an empty string if not present."""
        path = data_dict.get("path", "")

        for split in ["train", "val", "test"]:
@@ -386,6 +400,7 @@ class CometLogger:
        return data_dict

    def on_pretrain_routine_end(self, paths):
+        """Called at the end of pretraining routine to handle paths if training is not being resumed."""
        if self.opt.resume:
            return

@@ -398,20 +413,25 @@ class CometLogger:
        return

    def on_train_start(self):
+        """Logs hyperparameters at the start of training."""
        self.log_parameters(self.hyp)

    def on_train_epoch_start(self):
+        """Called at the start of each training epoch."""
        return

    def on_train_epoch_end(self, epoch):
+        """Updates the current epoch in the experiment tracking at the end of each epoch."""
        self.experiment.curr_epoch = epoch

        return

    def on_train_batch_start(self):
+        """Called at the start of each training batch."""
        return

    def on_train_batch_end(self, log_dict, step):
+        """Callback function that updates and logs metrics at the end of each training batch if conditions are met."""
        self.experiment.curr_step = step
        if self.log_batch_metrics and (step % self.comet_log_batch_interval == 0):
            self.log_metrics(log_dict, step=step)
@@ -419,6 +439,7 @@ class CometLogger:
        return

    def on_train_end(self, files, save_dir, last, best, epoch, results):
+        """Logs metadata and optionally saves model files at the end of training."""
        if self.comet_log_predictions:
            curr_epoch = self.experiment.curr_epoch
            self.experiment.log_asset_data(self.metadata_dict, "image-metadata.json", epoch=curr_epoch)
@@ -446,12 +467,15 @@ class CometLogger:
        self.finish_run()

    def on_val_start(self):
+        """Called at the start of validation, currently a placeholder with no functionality."""
        return

    def on_val_batch_start(self):
+        """Placeholder called at the start of a validation batch with no current functionality."""
        return

    def on_val_batch_end(self, batch_i, images, targets, paths, shapes, outputs):
+        """Callback executed at the end of a validation batch, conditionally logs predictions to Comet ML."""
        if not (self.comet_log_predictions and ((batch_i + 1) % self.comet_log_prediction_interval == 0)):
            return

@@ -470,6 +494,7 @@ class CometLogger:
        return

    def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
+        """Logs per-class metrics to Comet.ml after validation if enabled and more than one class exists."""
        if self.comet_log_per_class_metrics and self.num_classes > 1:
            for i, c in enumerate(ap_class):
                class_name = self.class_names[c]
@@ -504,14 +529,18 @@ class CometLogger:
            )

    def on_fit_epoch_end(self, result, epoch):
+        """Logs metrics at the end of each training epoch."""
        self.log_metrics(result, epoch=epoch)

    def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
+        """Callback to save model checkpoints periodically if conditions are met."""
        if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
            self.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)

    def on_params_update(self, params):
+        """Logs updated parameters during training."""
        self.log_parameters(params)

    def finish_run(self):
+        """Ends the current experiment and logs its completion."""
        self.experiment.end()
@@ -17,6 +17,7 @@ COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv("COMET_DEFAULT_CHECKPOINT_FILENAME


 def download_model_checkpoint(opt, experiment):
+    """Downloads YOLOv5 model checkpoint from Comet ML experiment, updating `opt.weights` with download path."""
    model_dir = f"{opt.project}/{experiment.name}"
    os.makedirs(model_dir, exist_ok=True)

@@ -25,6 +25,9 @@ COMET_PROJECT_NAME = config.get_string(os.getenv("COMET_PROJECT_NAME"), "comet.p


 def get_args(known=False):
+    """Parses command-line arguments for YOLOv5 training, supporting configuration of weights, data paths,
+    hyperparameters, and more.
+    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
    parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
@@ -83,6 +86,7 @@ def get_args(known=False):


 def run(parameters, opt):
+    """Executes YOLOv5 training with given hyperparameters and options, setting up device and training directories."""
    hyp_dict = {k: v for k, v in parameters.items() if k not in ["epochs", "batch_size"]}

    opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve))
@@ -152,6 +152,7 @@ class WandbLogger:
        LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")

    def val_one_image(self, pred, predn, path, names, im):
+        """Evaluates model prediction for a single image, returning metrics and visualizations."""
        pass

    def log(self, log_dict):