wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/integration/yolov8/__init__.py

@@ -0,0 +1,7 @@
+"""Tools for integrating `wandb` with [`ultralytics YOLOv8`](https://docs.ultralytics.com/).
+
+YOLOv8 is a computer vision framework for training and deploying YOLOv8 models.
+"""
+__all__ = ("add_callbacks",)
+
+from .yolov8 import add_callbacks

wandb/integration/yolov8/yolov8.py

@@ -0,0 +1,250 @@
+from typing import Any, Callable, Dict, List, Optional
+
+from ultralytics.yolo.engine.model import YOLO
+from ultralytics.yolo.engine.trainer import BaseTrainer
+from ultralytics.yolo.utils import RANK
+from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
+from ultralytics.yolo.v8.classify.train import ClassificationTrainer
+
+import wandb
+from wandb.sdk.lib import telemetry
+
+
+class WandbCallback:
+    """An internal YOLO model wrapper that tracks metrics, and logs models to Weights & Biases.
+
+    Usage:
+    ```python
+    from wandb.integration.yolov8.yolov8 import WandbCallback
+    model = YOLO("yolov8n.pt")
+    wandb_logger = WandbCallback(model,)
+    for event, callback_fn in wandb_logger.callbacks.items():
+        model.add_callback(event, callback_fn)
+    ```
+    """
+
+    def __init__(
+        self,
+        yolo: YOLO,
+        run_name: Optional[str] = None,
+        project: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        resume: Optional[str] = None,
+        **kwargs: Optional[Any],
+    ) -> None:
+        """A utility class to manage wandb run and various callbacks for the ultralytics YOLOv8 framework.
+
+        Args:
+            yolo: A YOLOv8 model that's inherited from `:class:ultralytics.yolo.engine.model.YOLO`
+            run_name, str: The name of the Weights & Biases run, defaults to an auto generated run_name if `trainer.args.name` is not defined.
+            project, str: The name of the Weights & Biases project, defaults to `"YOLOv8"` if `trainer.args.project` is not defined.
+            tags, List[str]: A list of tags to be added to the Weights & Biases run, defaults to `["YOLOv8"]`.
+            resume, str: Whether to resume a previous run on Weights & Biases, defaults to `None`.
+            **kwargs: Additional arguments to be passed to `wandb.init()`.
+        """
+        self.yolo = yolo
+        self.run_name = run_name
+        self.project = project
+        self.tags = tags
+        self.resume = resume
+        self.kwargs = kwargs
+
+    def on_pretrain_routine_start(self, trainer: BaseTrainer) -> None:
+        """Starts a new wandb run to track the training process and log to Weights & Biases.
+
+        Args:
+            trainer: A task trainer that's inherited from `:class:ultralytics.yolo.engine.trainer.BaseTrainer`
+                    that contains the model training and optimization routine.
+        """
+        if wandb.run is None:
+            self.run = wandb.init(
+                name=self.run_name if self.run_name else trainer.args.name,
+                project=self.project
+                if self.project
+                else trainer.args.project or "YOLOv8",
+                tags=self.tags if self.tags else ["YOLOv8"],
+                config=vars(trainer.args),
+                resume=self.resume if self.resume else None,
+                **self.kwargs,
+            )
+        else:
+            self.run = wandb.run
+        self.run.define_metric("epoch", hidden=True)
+        self.run.define_metric(
+            "train/*", step_metric="epoch", step_sync=True, summary="min"
+        )
+
+        self.run.define_metric(
+            "val/*", step_metric="epoch", step_sync=True, summary="min"
+        )
+
+        self.run.define_metric(
+            "metrics/*", step_metric="epoch", step_sync=True, summary="max"
+        )
+        self.run.define_metric(
+            "lr/*", step_metric="epoch", step_sync=True, summary="last"
+        )
+
+        with telemetry.context(run=wandb.run) as tel:
+            tel.feature.ultralytics_yolov8 = True
+
+    def on_pretrain_routine_end(self, trainer: BaseTrainer) -> None:
+        self.run.summary.update(
+            {
+                "model/parameters": get_num_params(trainer.model),
+                "model/GFLOPs": round(get_flops(trainer.model), 3),
+            }
+        )
+
+    def on_train_epoch_start(self, trainer: BaseTrainer) -> None:
+        """On train epoch start we only log epoch number to the Weights & Biases run."""
+        # We log the epoch number here to commit the previous step,
+        self.run.log({"epoch": trainer.epoch + 1})
+
+    def on_train_epoch_end(self, trainer: BaseTrainer) -> None:
+        """On train epoch end we log all the metrics to the Weights & Biases run."""
+        self.run.log(
+            {
+                **trainer.metrics,
+                **trainer.label_loss_items(trainer.tloss, prefix="train"),
+                **trainer.lr,
+            },
+        )
+        # Currently only the detection and segmentation trainers save images to the save_dir
+        if not isinstance(trainer, ClassificationTrainer):
+            self.run.log(
+                {
+                    "train_batch_images": [
+                        wandb.Image(str(image_path), caption=image_path.stem)
+                        for image_path in trainer.save_dir.glob("train_batch*.jpg")
+                    ]
+                }
+            )
+
+    def on_fit_epoch_end(self, trainer: BaseTrainer) -> None:
+        """On fit epoch end we log all the best metrics and model detail to Weights & Biases run summary."""
+        if trainer.epoch == 0:
+            speeds = [
+                trainer.validator.speed.get(
+                    key,
+                )
+                for key in (1, "inference")
+            ]
+            speed = speeds[0] if speeds[0] else speeds[1]
+            if speed:
+                self.run.summary.update(
+                    {
+                        "model/speed(ms/img)": round(speed, 3),
+                    }
+                )
+        if trainer.best_fitness == trainer.fitness:
+            self.run.summary.update(
+                {
+                    "best/epoch": trainer.epoch + 1,
+                    **{f"best/{key}": val for key, val in trainer.metrics.items()},
+                }
+            )
+
+    def on_train_end(self, trainer: BaseTrainer) -> None:
+        """On train end we log all the media, including plots, images and best model artifact to Weights & Biases."""
+        # Currently only the detection and segmentation trainers save images to the save_dir
+        if not isinstance(trainer, ClassificationTrainer):
+            self.run.log(
+                {
+                    "plots": [
+                        wandb.Image(str(image_path), caption=image_path.stem)
+                        for image_path in trainer.save_dir.glob("*.png")
+                    ],
+                    "val_images": [
+                        wandb.Image(str(image_path), caption=image_path.stem)
+                        for image_path in trainer.validator.save_dir.glob("val*.jpg")
+                    ],
+                },
+            )
+
+        if trainer.best.exists():
+            self.run.log_artifact(
+                str(trainer.best),
+                type="model",
+                name=f"{self.run.name}_{trainer.args.task}.pt",
+                aliases=["best", f"epoch_{trainer.epoch + 1}"],
+            )
+
+    def on_model_save(self, trainer: BaseTrainer) -> None:
+        """On model save we log the model as an artifact to Weights & Biases."""
+        self.run.log_artifact(
+            str(trainer.last),
+            type="model",
+            name=f"{self.run.name}_{trainer.args.task}.pt",
+            aliases=["last", f"epoch_{trainer.epoch + 1}"],
+        )
+
+    def teardown(self, _trainer: BaseTrainer) -> None:
+        """On teardown, we finish the Weights & Biases run and set it to None."""
+        self.run.finish()
+        self.run = None
+
+    @property
+    def callbacks(
+        self,
+    ) -> Dict[str, Callable]:
+        """Property contains all the relevant callbacks to add to the YOLO model for the Weights & Biases logging."""
+        return {
+            "on_pretrain_routine_start": self.on_pretrain_routine_start,
+            "on_pretrain_routine_end": self.on_pretrain_routine_end,
+            "on_train_epoch_start": self.on_train_epoch_start,
+            "on_train_epoch_end": self.on_train_epoch_end,
+            "on_fit_epoch_end": self.on_fit_epoch_end,
+            "on_train_end": self.on_train_end,
+            "on_model_save": self.on_model_save,
+            "teardown": self.teardown,
+        }
+
+
+def add_callbacks(
+    yolo: YOLO,
+    run_name: Optional[str] = None,
+    project: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    resume: Optional[str] = None,
+    **kwargs: Optional[Any],
+) -> YOLO:
+    """A YOLO model wrapper that tracks metrics, and logs models to Weights & Biases.
+
+    Args:
+        yolo: A YOLOv8 model that's inherited from `:class:ultralytics.yolo.engine.model.YOLO`
+        run_name, str: The name of the Weights & Biases run, defaults to an auto generated name if `trainer.args.name` is not defined.
+        project, str: The name of the Weights & Biases project, defaults to `"YOLOv8"` if `trainer.args.project` is not defined.
+        tags, List[str]: A list of tags to be added to the Weights & Biases run, defaults to `["YOLOv8"]`.
+        resume, str: Whether to resume a previous run on Weights & Biases, defaults to `None`.
+        **kwargs: Additional arguments to be passed to `wandb.init()`.
+
+    Usage:
+    ```python
+    from wandb.integration.yolov8 import add_callbacks as add_wandb_callbacks
+    model = YOLO("yolov8n.pt")
+    add_wandb_callbacks(model,)
+    model.train(data="coco128.yaml", epochs=3, imgsz=640,)
+    ```
+    """
+    wandb.termwarn(
+        """The wandb callback is currently in beta and is subject to change based on updates to `ultralytics yolov8`.
+        The callback is tested and supported for ultralytics v8.0.43 and above.
+        Please report any issues to https://github.com/wandb/wandb/issues with the tag `yolov8`.
+        """,
+        repeat=False,
+    )
+
+    if RANK in [-1, 0]:
+        wandb_logger = WandbCallback(
+            yolo, run_name=run_name, project=project, tags=tags, resume=resume, **kwargs
+        )
+        for event, callback_fn in wandb_logger.callbacks.items():
+            yolo.add_callback(event, callback_fn)
+        return yolo
+    else:
+        wandb.termerror(
+            "The RANK of the process to add the callbacks was neither 0 or -1."
+            "No Weights & Biases callbacks were added to this instance of the YOLO model."
+        )
+    return yolo

wandb/jupyter.py

@@ -5,11 +5,13 @@ import re
 import shutil
 import sys
 from base64 import b64encode
+from typing import Dict
 
 import requests
 from requests.compat import urljoin
 
 import wandb
+import wandb.util
 from wandb.sdk.lib import filesystem
 
 try:
@@ -42,7 +44,7 @@ __IFrame = None
 
 
 def maybe_display():
-    """Display a run if the user added cell magic and we have run"""
+    """Display a run if the user added cell magic and we have run."""
     if __IFrame is not None:
         return __IFrame.maybe_display()
     return False
@@ -177,10 +179,10 @@ def notebook_metadata_from_jupyter_servers_and_kernel_id():
     return None
 
 
-def notebook_metadata(silent):
-    """Attempts to query jupyter for the path and name of the notebook file.
+def notebook_metadata(silent) -> Dict[str, str]:
+    """Attempt to query jupyter for the path and name of the notebook file.
 
-    This can handle many different jupyter environments, specifically:
+    This can handle different jupyter environments, specifically:
 
     1. Colab
     2. Kaggle
@@ -193,30 +195,21 @@ def notebook_metadata(silent):
         "the WANDB_NOTEBOOK_NAME environment variable to enable code saving."
     )
     try:
-        # In colab we can request the most recent contents
+        jupyter_metadata = notebook_metadata_from_jupyter_servers_and_kernel_id()
+
+        # Colab:
+        # request the most recent contents
         ipynb = attempt_colab_load_ipynb()
-        if ipynb:
-            nb_name = ipynb["metadata"]["colab"]["name"]
-            if ".ipynb" not in nb_name:
-                nb_name += ".ipynb"
-            ret = {
+        if ipynb is not None and jupyter_metadata is not None:
+            return {
                 "root": "/content",
-                "path": nb_name,
-                "name": nb_name,
+                "path": jupyter_metadata["path"],
+                "name": jupyter_metadata["name"],
             }
 
-            try:
-                jupyter_metadata = (
-                    notebook_metadata_from_jupyter_servers_and_kernel_id()
-                )
-            except RuntimeError:
-                pass
-            else:
-                ret["path"] = jupyter_metadata["path"]
-            return ret
-
+        # Kaggle:
         if wandb.util._is_kaggle():
-            # In kaggle we can request the most recent contents
+            # request the most recent contents
             ipynb = attempt_kaggle_load_ipynb()
             if ipynb:
                 return {
@@ -225,7 +218,6 @@ def notebook_metadata(silent):
                     "name": ipynb["metadata"]["name"],
                 }
 
-        jupyter_metadata = notebook_metadata_from_jupyter_servers_and_kernel_id()
         if jupyter_metadata:
             return jupyter_metadata
         if not silent:
@@ -241,8 +233,10 @@ def notebook_metadata(silent):
 
 
 def jupyter_servers_and_kernel_id():
-    """Returns a list of servers and the current kernel_id so we can query for
-    the name of the notebook"""
+    """Return a list of servers and the current kernel_id.
+
+    Used to query for the name of the notebook.
+    """
     try:
         import ipykernel
 
@@ -286,7 +280,7 @@ def attempt_kaggle_load_ipynb():
 
 
 def attempt_colab_login(app_url):
-    """This renders an iframe to wandb in the hopes it posts back an api key"""
+    """This renders an iframe to wandb in the hopes it posts back an api key."""
     from google.colab import output
     from google.colab._message import MessageError
     from IPython import display
@@ -321,7 +315,7 @@ def attempt_colab_login(app_url):
             });
             })
         });
-    """  # noqa: E501
+    """
             % app_url.replace("http:", "https:")
         )
     )
@@ -401,12 +395,14 @@ class Notebook:
         # TODO: likely only save if the code has changed
         colab_ipynb = attempt_colab_load_ipynb()
         if colab_ipynb:
-            nb_name = (
-                colab_ipynb.get("metadata", {})
-                .get("colab", {})
-                .get("name", "colab.ipynb")
-            )
-            if ".ipynb" not in nb_name:
+            try:
+                jupyter_metadata = (
+                    notebook_metadata_from_jupyter_servers_and_kernel_id()
+                )
+                nb_name = jupyter_metadata["name"]
+            except Exception:
+                nb_name = "colab.ipynb"
+            if not nb_name.endswith(".ipynb"):
                 nb_name += ".ipynb"
             with open(
                 os.path.join(
@@ -434,7 +430,7 @@ class Notebook:
         return False
 
     def save_history(self):
-        """This saves all cell executions in the current session as a new notebook"""
+        """This saves all cell executions in the current session as a new notebook."""
         try:
             from nbformat import v4, validator, write
         except ImportError:

wandb/lightgbm/__init__.py

@@ -1,5 +1,4 @@
-"""
-Compatibility lightgbm module.
+"""Compatibility lightgbm module.
 
 In the future use:
     from wandb.integration.lightgbm import wandb_callback

wandb/old/settings.py

@@ -43,7 +43,7 @@ class Settings:
                     raise
 
     def set(self, section, key, value, globally=False, persist=False) -> None:
-        """Persists settings to disk if persist = True"""
+        """Persist settings to disk if persist = True"""
 
         def write_setting(settings, settings_path, persist):
             if not settings.has_section(section):
@@ -109,7 +109,7 @@ class Settings:
         config_dir = os.environ.get(
             env.CONFIG_DIR, os.path.join(os.path.expanduser("~"), ".config", "wandb")
         )
-        filesystem.mkdir_exists_ok(config_dir)
+        os.makedirs(config_dir, exist_ok=True)
         return os.path.join(config_dir, "settings")
 
     @staticmethod

wandb/plot/bar.py

@@ -4,8 +4,7 @@ import wandb
 
 
 def bar(table: wandb.Table, label: str, value: str, title: Optional[str] = None):
-    """
-    Construct a bar plot.
+    """Construct a bar plot.
 
     Arguments:
         table (wandb.Table): Table of data.

wandb/plot/confusion_matrix.py

@@ -13,8 +13,7 @@ def confusion_matrix(
     class_names: Optional[Sequence[str]] = None,
     title: Optional[str] = None,
 ):
-    """
-    Computes a multi-run confusion matrix.
+    """Compute a multi-run confusion matrix.
 
     Arguments:
         probs (2-d arr): Shape [n_examples, n_classes]
@@ -35,7 +34,6 @@ def confusion_matrix(
         wandb.log({'confusion_matrix': wandb.plot.confusion_matrix(probs, y_true=y_true, class_names=labels)})
         ```
     """
-
     np = util.get_module(
         "numpy",
         required="confusion matrix requires the numpy library, install with `pip install numpy`",

wandb/plot/histogram.py

@@ -4,8 +4,7 @@ import wandb
 
 
 def histogram(table: wandb.Table, value: str, title: Optional[str] = None):
-    """
-    Construct a histogram plot.
+    """Construct a histogram plot.
 
     Arguments:
         table (wandb.Table): Table of data.

wandb/plot/line.py

@@ -10,8 +10,7 @@ def line(
     stroke: Optional[str] = None,
     title: Optional[str] = None,
 ):
-    """
-    Construct a line plot.
+    """Construct a line plot.
 
     Arguments:
         table (wandb.Table): Table of data.

wandb/plot/line_series.py

@@ -11,8 +11,7 @@ def line_series(
     title: t.Optional[str] = None,
     xname: t.Optional[str] = None,
 ):
-    """
-    Construct a line series plot.
+    """Construct a line series plot.
 
     Arguments:
         xs (array of arrays, or array): Array of arrays of x values
@@ -44,10 +43,11 @@ def line_series(
         import wandb
 
         run = wandb.init()
-        xs = [[i for i in range(10)], [2*i for i in range(10)]]
+        xs = [[i for i in range(10)], [2 * i for i in range(10)]]
         ys = [[i for i in range(10)], [i**2 for i in range(10)]]
         run.log(
-            {'line-series-plot2': wandb.plot.line_series(xs, ys, title="title", xname="step")})
+            {"line-series-plot2": wandb.plot.line_series(xs, ys, title="title", xname="step")}
+        )
         run.finish()
         ```
     """

wandb/plot/pr_curve.py

@@ -11,32 +11,29 @@ def pr_curve(
     interp_size=21,
     title=None,
 ):
-    """
-    Computes the tradeoff between precision and recall for different thresholds.
-    A high area under the curve represents both high recall and high precision,
-    where high precision relates to a low false positive rate, and high recall
-    relates to a low false negative rate. High scores for both show that the
-    classifier is returning accurate results (high precision), and
-    returning a majority of all positive results (high recall).
-    PR curve is useful when the classes are very imbalanced.
+    """Compute the tradeoff between precision and recall for different thresholds.
+
+    A high area under the curve represents both high recall and high precision, where
+    high precision relates to a low false positive rate, and high recall relates to a
+    low false negative rate. High scores for both show that the classifier is returning
+    accurate results (high precision), and returning a majority of all positive results
+    (high recall). PR curve is useful when the classes are very imbalanced.
 
     Arguments:
-        y_true (arr): true sparse labels
-        y_probas (arr): Target scores, can either be probability estimates, confidence
-                         values, or non-thresholded measure of decisions.
-                         shape: (*y_true.shape, num_classes)
-        labels (list): Named labels for target variable (y). Makes plots easier to
-                        read by replacing target values with corresponding index.
-                        For example labels = ['dog', 'cat', 'owl'] all 0s are
-                        replaced by 'dog', 1s by 'cat'.
+        y_true (arr): true sparse labels y_probas (arr): Target scores, can either be
+            probability estimates, confidence values, or non-thresholded measure of
+            decisions. shape: (*y_true.shape, num_classes)
+        labels (list): Named labels for target variable (y). Makes plots easier to read
+            by replacing target values with corresponding index. For example labels =
+            ['dog', 'cat', 'owl'] all 0s are replaced by 'dog', 1s by 'cat'.
         classes_to_plot (list): unique values of y_true to include in the plot
         interp_size (int): the recall values will be fixed to `interp_size` points
-                            uniform on [0, 1] and the precision will be interpolated for
-                            these recall values.
+            uniform on [0, 1] and the precision will be interpolated for these recall
+            values.
 
     Returns:
-        Nothing. To see plots, go to your W&B run page then expand the 'media' tab
-        under 'auto visualizations'.
+        Nothing. To see plots, go to your W&B run page then expand the 'media' tab under
+        'auto visualizations'.
 
     Example:
         ```

wandb/plot/roc_curve.py

@@ -6,9 +6,7 @@ from wandb.plots.utils import test_missing, test_types
 def roc_curve(
     y_true=None, y_probas=None, labels=None, classes_to_plot=None, title=None
 ):
-    """
-    Calculates receiver operating characteristic scores and visualizes them as the
-    ROC curve.
+    """Calculate and visualize receiver operating characteristic (ROC) scores.
 
     Arguments:
         y_true (arr): true sparse labels

wandb/plot/scatter.py

@@ -2,8 +2,7 @@ import wandb
 
 
 def scatter(table, x, y, title=None):
-    """
-    Construct a scatter plot.
+    """Construct a scatter plot.
 
     Arguments:
         table (wandb.Table): Table of data.