wandb 0.21.2__py3-none-macosx_12_0_x86_64.whl

wandb/integration/keras/callbacks/__init__.py

@@ -0,0 +1,5 @@
+__all__ = ("WandbMetricsLogger", "WandbModelCheckpoint", "WandbEvalCallback")
+
+from .metrics_logger import WandbMetricsLogger
+from .model_checkpoint import WandbModelCheckpoint
+from .tables_builder import WandbEvalCallback

wandb/integration/keras/callbacks/metrics_logger.py

@@ -0,0 +1,129 @@
+from typing import Any, Dict, Literal, Optional, Union
+
+import tensorflow as tf  # type: ignore
+from tensorflow.keras import callbacks
+
+import wandb
+from wandb.integration.keras.keras import patch_tf_keras
+from wandb.sdk.lib import telemetry
+
+LogStrategy = Literal["epoch", "batch"]
+
+
+patch_tf_keras()
+
+
+class WandbMetricsLogger(callbacks.Callback):
+    """Logger that sends system metrics to W&B.
+
+    `WandbMetricsLogger` automatically logs the `logs` dictionary that callback methods
+    take as argument to wandb.
+
+    This callback automatically logs the following to a W&B run page:
+    * system (CPU/GPU/TPU) metrics,
+    * train and validation metrics defined in `model.compile`,
+    * learning rate (both for a fixed value or a learning rate scheduler)
+
+    Notes:
+    If you resume training by passing `initial_epoch` to `model.fit` and you are using a
+    learning rate scheduler, make sure to pass `initial_global_step` to
+    `WandbMetricsLogger`. The `initial_global_step` is `step_size * initial_step`, where
+    `step_size` is number of training steps per epoch. `step_size` can be calculated as
+    the product of the cardinality of the training dataset and the batch size.
+
+    Args:
+        log_freq: ("epoch", "batch", or int) if "epoch", logs metrics
+            at the end of each epoch. If "batch", logs metrics at the end
+            of each batch. If an integer, logs metrics at the end of that
+            many batches. Defaults to "epoch".
+        initial_global_step: (int) Use this argument to correctly log the
+            learning rate when you resume training from some `initial_epoch`,
+            and a learning rate scheduler is used. This can be computed as
+            `step_size * initial_step`. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        log_freq: Union[LogStrategy, int] = "epoch",
+        initial_global_step: int = 0,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        if wandb.run is None:
+            raise wandb.Error(
+                "You must call `wandb.init()` before WandbMetricsLogger()"
+            )
+
+        with telemetry.context(run=wandb.run) as tel:
+            tel.feature.keras_metrics_logger = True
+
+        if log_freq == "batch":
+            log_freq = 1
+
+        self.logging_batch_wise = isinstance(log_freq, int)
+        self.log_freq: Any = log_freq if self.logging_batch_wise else None
+        self.global_batch = 0
+        self.global_step = initial_global_step
+
+        if self.logging_batch_wise:
+            # define custom x-axis for batch logging.
+            wandb.define_metric("batch/batch_step")
+            # set all batch metrics to be logged against batch_step.
+            wandb.define_metric("batch/*", step_metric="batch/batch_step")
+        else:
+            # define custom x-axis for epoch-wise logging.
+            wandb.define_metric("epoch/epoch")
+            # set all epoch-wise metrics to be logged against epoch.
+            wandb.define_metric("epoch/*", step_metric="epoch/epoch")
+
+    def _get_lr(self) -> Union[float, None]:
+        if isinstance(
+            self.model.optimizer.learning_rate,
+            (tf.Variable, tf.Tensor),
+        ) or (
+            hasattr(self.model.optimizer.learning_rate, "shape")
+            and self.model.optimizer.learning_rate.shape == ()
+        ):
+            return float(self.model.optimizer.learning_rate.numpy().item())
+        try:
+            return float(
+                self.model.optimizer.learning_rate(step=self.global_step).numpy().item()
+            )
+        except Exception as e:
+            wandb.termerror(f"Unable to log learning rate: {e}", repeat=False)
+            return None
+
+    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None) -> None:
+        """Called at the end of an epoch."""
+        logs = dict() if logs is None else {f"epoch/{k}": v for k, v in logs.items()}
+
+        logs["epoch/epoch"] = epoch
+
+        lr = self._get_lr()
+        if lr is not None:
+            logs["epoch/learning_rate"] = lr
+
+        wandb.log(logs)
+
+    def on_batch_end(self, batch: int, logs: Optional[Dict[str, Any]] = None) -> None:
+        self.global_step += 1
+        """An alias for `on_train_batch_end` for backwards compatibility."""
+        if self.logging_batch_wise and batch % self.log_freq == 0:
+            logs = {f"batch/{k}": v for k, v in logs.items()} if logs else {}
+            logs["batch/batch_step"] = self.global_batch
+
+            lr = self._get_lr()
+            if lr is not None:
+                logs["batch/learning_rate"] = lr
+
+            wandb.log(logs)
+
+            self.global_batch += self.log_freq
+
+    def on_train_batch_end(
+        self, batch: int, logs: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Called at the end of a training batch in `fit` methods."""
+        self.on_batch_end(batch, logs if logs else {})

wandb/integration/keras/callbacks/model_checkpoint.py

@@ -0,0 +1,188 @@
+import os
+import string
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import tensorflow as tf  # type: ignore
+from tensorflow.keras import callbacks  # type: ignore
+
+import wandb
+from wandb.sdk.lib import telemetry
+from wandb.sdk.lib.paths import StrPath
+
+from ..keras import patch_tf_keras
+
+Mode = Literal["auto", "min", "max"]
+SaveStrategy = Literal["epoch"]
+
+patch_tf_keras()
+
+
+class WandbModelCheckpoint(callbacks.ModelCheckpoint):
+    """A checkpoint that periodically saves a Keras model or model weights.
+
+    Saved weights are uploaded to W&B as a `wandb.Artifact`.
+
+    Since this callback is subclassed from `tf.keras.callbacks.ModelCheckpoint`, the
+    checkpointing logic is taken care of by the parent callback. You can learn more
+    here: https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint
+
+    This callback is to be used in conjunction with training using `model.fit()` to save
+    a model or weights (in a checkpoint file) at some interval. The model checkpoints
+    will be logged as W&B Artifacts. You can learn more here:
+    https://docs.wandb.ai/guides/artifacts
+
+    This callback provides the following features:
+        - Save the model that has achieved "best performance" based on "monitor".
+        - Save the model at the end of every epoch regardless of the performance.
+        - Save the model at the end of epoch or after a fixed number of training batches.
+        - Save only model weights, or save the whole model.
+        - Save the model either in SavedModel format or in `.h5` format.
+
+    Args:
+        filepath: (Union[str, os.PathLike]) path to save the model file. `filepath`
+            can contain named formatting options, which will be filled by the value
+            of `epoch` and keys in `logs` (passed in `on_epoch_end`). For example:
+            if `filepath` is `model-{epoch:02d}-{val_loss:.2f}`, then the
+            model checkpoints will be saved with the epoch number and the
+            validation loss in the filename.
+        monitor: (str) The metric name to monitor. Default to "val_loss".
+        verbose: (int) Verbosity mode, 0 or 1. Mode 0 is silent, and mode 1
+            displays messages when the callback takes an action.
+        save_best_only: (bool) if `save_best_only=True`, it only saves when the model
+            is considered the "best" and the latest best model according to the
+            quantity monitored will not be overwritten. If `filepath` doesn't contain
+            formatting options like `{epoch}` then `filepath` will be overwritten by
+            each new better model locally. The model logged as an artifact will still be
+            associated with the correct `monitor`.  Artifacts will be uploaded
+            continuously and versioned separately as a new best model is found.
+        save_weights_only: (bool) if True, then only the model's weights will be saved.
+        mode: (Mode) one of {'auto', 'min', 'max'}. For `val_acc`, this should be `max`,
+            for `val_loss` this should be `min`, etc.
+        save_freq: (Union[SaveStrategy, int]) `epoch` or integer. When using `'epoch'`,
+            the callback saves the model after each epoch. When using an integer, the
+            callback saves the model at end of this many batches.
+            Note that when monitoring validation metrics such as `val_acc` or `val_loss`,
+            save_freq must be set to "epoch" as those metrics are only available at the
+            end of an epoch.
+        initial_value_threshold: (Optional[float]) Floating point initial "best" value of the metric
+            to be monitored.
+    """
+
+    def __init__(
+        self,
+        filepath: StrPath,
+        monitor: str = "val_loss",
+        verbose: int = 0,
+        save_best_only: bool = False,
+        save_weights_only: bool = False,
+        mode: Mode = "auto",
+        save_freq: Union[SaveStrategy, int] = "epoch",
+        initial_value_threshold: Optional[float] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            filepath=filepath,
+            monitor=monitor,
+            verbose=verbose,
+            save_best_only=save_best_only,
+            save_weights_only=save_weights_only,
+            mode=mode,
+            save_freq=save_freq,
+            initial_value_threshold=initial_value_threshold,
+            **kwargs,
+        )
+        if wandb.run is None:
+            raise wandb.Error(
+                "You must call `wandb.init()` before `WandbModelCheckpoint()`"
+            )
+        with telemetry.context(run=wandb.run) as tel:
+            tel.feature.keras_model_checkpoint = True
+
+        self.save_weights_only = save_weights_only
+
+        # User-friendly warning when trying to save the best model.
+        if self.save_best_only:
+            self._check_filepath()
+
+        self._is_old_tf_keras_version: Optional[bool] = None
+
+    def on_train_batch_end(
+        self, batch: int, logs: Optional[Dict[str, float]] = None
+    ) -> None:
+        if self._should_save_on_batch(batch):
+            if self.is_old_tf_keras_version:
+                # Save the model and get filepath
+                self._save_model(epoch=self._current_epoch, logs=logs)
+                filepath = self._get_file_path(epoch=self._current_epoch, logs=logs)
+            else:
+                # Save the model and get filepath
+                self._save_model(epoch=self._current_epoch, batch=batch, logs=logs)
+                filepath = self._get_file_path(
+                    epoch=self._current_epoch, batch=batch, logs=logs
+                )
+            # Log the model as artifact
+            aliases = ["latest", f"epoch_{self._current_epoch}_batch_{batch}"]
+            self._log_ckpt_as_artifact(filepath, aliases=aliases)
+
+    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, float]] = None) -> None:
+        super().on_epoch_end(epoch, logs)
+        # Check if model checkpoint is created at the end of epoch.
+        if self.save_freq == "epoch":
+            # Get filepath where the model checkpoint is saved.
+            if self.is_old_tf_keras_version:
+                filepath = self._get_file_path(epoch=epoch, logs=logs)
+            else:
+                filepath = self._get_file_path(epoch=epoch, batch=None, logs=logs)
+            # Log the model as artifact
+            aliases = ["latest", f"epoch_{epoch}"]
+            self._log_ckpt_as_artifact(filepath, aliases=aliases)
+
+    def _log_ckpt_as_artifact(
+        self, filepath: str, aliases: Optional[List[str]] = None
+    ) -> None:
+        """Log model checkpoint as  W&B Artifact."""
+        try:
+            assert wandb.run is not None
+            model_checkpoint_artifact = wandb.Artifact(
+                f"run_{wandb.run.id}_model", type="model"
+            )
+            if os.path.isfile(filepath):
+                model_checkpoint_artifact.add_file(filepath)
+            elif os.path.isdir(filepath):
+                model_checkpoint_artifact.add_dir(filepath)
+            else:
+                raise FileNotFoundError(f"No such file or directory {filepath}")
+            wandb.log_artifact(model_checkpoint_artifact, aliases=aliases or [])
+        except ValueError:
+            # This error occurs when `save_best_only=True` and the model
+            # checkpoint is not saved for that epoch/batch. Since TF/Keras
+            # is giving friendly log, we can avoid clustering the stdout.
+            pass
+
+    def _check_filepath(self) -> None:
+        placeholders = []
+        for tup in string.Formatter().parse(self.filepath):
+            if tup[1] is not None:
+                placeholders.append(tup[1])
+        if len(placeholders) == 0:
+            wandb.termwarn(
+                "When using `save_best_only`, ensure that the `filepath` argument "
+                "contains formatting placeholders like `{epoch:02d}` or `{batch:02d}`. "
+                "This ensures correct interpretation of the logged artifacts.",
+                repeat=False,
+            )
+
+    @property
+    def is_old_tf_keras_version(self) -> Optional[bool]:
+        if self._is_old_tf_keras_version is None:
+            from packaging.version import parse
+
+            try:
+                if parse(tf.keras.__version__) < parse("2.6.0"):
+                    self._is_old_tf_keras_version = True
+                else:
+                    self._is_old_tf_keras_version = False
+            except AttributeError:
+                self._is_old_tf_keras_version = False
+
+        return self._is_old_tf_keras_version

wandb/integration/keras/callbacks/tables_builder.py

@@ -0,0 +1,228 @@
+import abc
+from typing import Any, Dict, List, Optional
+
+from tensorflow.keras.callbacks import Callback  # type: ignore
+
+import wandb
+from wandb.sdk.lib import telemetry
+
+
+class WandbEvalCallback(Callback, abc.ABC):
+    """Abstract base class to build Keras callbacks for model prediction visualization.
+
+    You can build callbacks for visualizing model predictions `on_epoch_end`
+    that can be passed to `model.fit()` for classification, object detection,
+    segmentation, etc. tasks.
+
+    To use this, inherit from this base callback class and implement the
+    `add_ground_truth` and `add_model_prediction` methods.
+
+    The base class will take care of the following:
+    - Initialize `data_table` for logging the ground truth and
+        `pred_table` for predictions.
+    - The data uploaded to `data_table` is used as a reference for the
+        `pred_table`. This is to reduce the memory footprint. The `data_table_ref`
+        is a list that can be used to access the referenced data.
+        Check out the example below to see how it's done.
+    - Log the tables to W&B as W&B Artifacts.
+    - Each new `pred_table` is logged as a new version with aliases.
+
+    Example:
+        ```python
+        class WandbClfEvalCallback(WandbEvalCallback):
+            def __init__(self, validation_data, data_table_columns, pred_table_columns):
+                super().__init__(data_table_columns, pred_table_columns)
+
+                self.x = validation_data[0]
+                self.y = validation_data[1]
+
+            def add_ground_truth(self):
+                for idx, (image, label) in enumerate(zip(self.x, self.y)):
+                    self.data_table.add_data(idx, wandb.Image(image), label)
+
+            def add_model_predictions(self, epoch):
+                preds = self.model.predict(self.x, verbose=0)
+                preds = tf.argmax(preds, axis=-1)
+
+                data_table_ref = self.data_table_ref
+                table_idxs = data_table_ref.get_index()
+
+                for idx in table_idxs:
+                    pred = preds[idx]
+                    self.pred_table.add_data(
+                        epoch,
+                        data_table_ref.data[idx][0],
+                        data_table_ref.data[idx][1],
+                        data_table_ref.data[idx][2],
+                        pred,
+                    )
+
+
+        model.fit(
+            x,
+            y,
+            epochs=2,
+            validation_data=(x, y),
+            callbacks=[
+                WandbClfEvalCallback(
+                    validation_data=(x, y),
+                    data_table_columns=["idx", "image", "label"],
+                    pred_table_columns=["epoch", "idx", "image", "label", "pred"],
+                )
+            ],
+        )
+        ```
+
+    To have more fine-grained control, you can override the `on_train_begin` and
+    `on_epoch_end` methods. If you want to log the samples after N batched, you
+    can implement `on_train_batch_end` method.
+    """
+
+    def __init__(
+        self,
+        data_table_columns: List[str],
+        pred_table_columns: List[str],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        if wandb.run is None:
+            raise wandb.Error(
+                "You must call `wandb.init()` first before using this callback."
+            )
+
+        with telemetry.context(run=wandb.run) as tel:
+            tel.feature.keras_wandb_eval_callback = True
+
+        self.data_table_columns = data_table_columns
+        self.pred_table_columns = pred_table_columns
+
+    def on_train_begin(self, logs: Optional[Dict[str, float]] = None) -> None:
+        # Initialize the data_table
+        self.init_data_table(column_names=self.data_table_columns)
+        # Log the ground truth data
+        self.add_ground_truth(logs)
+        # Log the data_table as W&B Artifacts
+        self.log_data_table()
+
+    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, float]] = None) -> None:
+        # Initialize the pred_table
+        self.init_pred_table(column_names=self.pred_table_columns)
+        # Log the model prediction
+        self.add_model_predictions(epoch, logs)
+        # Log the pred_table as W&B Artifacts
+        self.log_pred_table()
+
+    @abc.abstractmethod
+    def add_ground_truth(self, logs: Optional[Dict[str, float]] = None) -> None:
+        """Add ground truth data to `data_table`.
+
+        Use this method to write the logic for adding validation/training data to
+        `data_table` initialized using `init_data_table` method.
+
+        Example:
+            ```python
+            for idx, data in enumerate(dataloader):
+                self.data_table.add_data(idx, data)
+            ```
+        This method is called once `on_train_begin` or equivalent hook.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__}.add_ground_truth")
+
+    @abc.abstractmethod
+    def add_model_predictions(
+        self, epoch: int, logs: Optional[Dict[str, float]] = None
+    ) -> None:
+        """Add a prediction from a model to `pred_table`.
+
+        Use this method to write the logic for adding model prediction for validation/
+        training data to `pred_table` initialized using `init_pred_table` method.
+
+        Example:
+            ```python
+            # Assuming the dataloader is not shuffling the samples.
+            for idx, data in enumerate(dataloader):
+                preds = model.predict(data)
+                self.pred_table.add_data(
+                    self.data_table_ref.data[idx][0],
+                    self.data_table_ref.data[idx][1],
+                    preds,
+                )
+            ```
+        This method is called `on_epoch_end` or equivalent hook.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__}.add_model_predictions")
+
+    def init_data_table(self, column_names: List[str]) -> None:
+        """Initialize the W&B Tables for validation data.
+
+        Call this method `on_train_begin` or equivalent hook. This is followed by adding
+        data to the table row or column wise.
+
+        Args:
+            column_names: (list) Column names for W&B Tables.
+        """
+        self.data_table = wandb.Table(columns=column_names, allow_mixed_types=True)
+
+    def init_pred_table(self, column_names: List[str]) -> None:
+        """Initialize the W&B Tables for model evaluation.
+
+        Call this method `on_epoch_end` or equivalent hook. This is followed by adding
+        data to the table row or column wise.
+
+        Args:
+            column_names: (list) Column names for W&B Tables.
+        """
+        self.pred_table = wandb.Table(columns=column_names)
+
+    def log_data_table(
+        self, name: str = "val", type: str = "dataset", table_name: str = "val_data"
+    ) -> None:
+        """Log the `data_table` as W&B artifact and call `use_artifact` on it.
+
+        This lets the evaluation table use the reference of already uploaded data
+        (images, text, scalar, etc.) without re-uploading.
+
+        Args:
+            name: (str) A human-readable name for this artifact, which is how you can
+                identify this artifact in the UI or reference it in use_artifact calls.
+                (default is 'val')
+            type: (str) The type of the artifact, which is used to organize and
+                differentiate artifacts. (default is 'dataset')
+            table_name: (str) The name of the table as will be displayed in the UI.
+                (default is 'val_data').
+        """
+        data_artifact = wandb.Artifact(name, type=type)
+        data_artifact.add(self.data_table, table_name)
+
+        # Calling `use_artifact` uploads the data to W&B.
+        assert wandb.run is not None
+        wandb.run.use_artifact(data_artifact)
+        data_artifact.wait()
+
+        # We get the reference table.
+        self.data_table_ref = data_artifact.get(table_name)
+
+    def log_pred_table(
+        self,
+        type: str = "evaluation",
+        table_name: str = "eval_data",
+        aliases: Optional[List[str]] = None,
+    ) -> None:
+        """Log the W&B Tables for model evaluation.
+
+        The table will be logged multiple times creating new version. Use this
+        to compare models at different intervals interactively.
+
+        Args:
+            type: (str) The type of the artifact, which is used to organize and
+                differentiate artifacts. (default is 'evaluation')
+            table_name: (str) The name of the table as will be displayed in the UI.
+                (default is 'eval_data')
+            aliases: (List[str]) List of aliases for the prediction table.
+        """
+        assert wandb.run is not None
+        pred_artifact = wandb.Artifact(f"run_{wandb.run.id}_pred", type=type)
+        pred_artifact.add(self.pred_table, table_name)
+        wandb.run.log_artifact(pred_artifact, aliases=aliases or ["latest"])