star

venv/lib/python3.9/site-packages/mlflow/tensorflow/autologging.py

@@ -0,0 +1,184 @@
+from typing import Union
+
+import numpy as np
+import tensorflow
+from tensorflow.keras.callbacks import TensorBoard
+
+from mlflow.utils.autologging_utils import (
+    INPUT_EXAMPLE_SAMPLE_ROWS,
+    ExceptionSafeClass,
+)
+
+
+class _TensorBoard(TensorBoard, metaclass=ExceptionSafeClass):
+    pass
+
+
+def _extract_input_example_from_tensor_or_ndarray(
+    input_features: Union[tensorflow.Tensor, np.ndarray],
+) -> np.ndarray:
+    """
+    Extracts first `INPUT_EXAMPLE_SAMPLE_ROWS` from the next_input, which can either be of
+    numpy array or tensor type.
+
+    Args:
+        input_features: an input of type `np.ndarray` or `tensorflow.Tensor`
+
+    Returns:
+        A slice (of limit `INPUT_EXAMPLE_SAMPLE_ROWS`)  of the input of type `np.ndarray`.
+        Returns `None` if the type of `input_features` is unsupported.
+
+    Examples
+    --------
+    when next_input is nd.array:
+    >>> input_data = np.array([1, 2, 3, 4, 5, 6, 7, 8])
+    >>> _extract_input_example_from_tensor_or_ndarray(input_data)
+    array([1, 2, 3, 4, 5])
+
+
+    when next_input is tensorflow.Tensor:
+    >>> input_data = tensorflow.convert_to_tensor([1, 2, 3, 4, 5, 6])
+    >>> _extract_input_example_from_tensor_or_ndarray(input_data)
+    array([1, 2, 3, 4, 5])
+    """
+
+    input_feature_slice = None
+    if isinstance(input_features, tensorflow.Tensor):
+        input_feature_slice = input_features.numpy()[0:INPUT_EXAMPLE_SAMPLE_ROWS]
+    elif isinstance(input_features, np.ndarray):
+        input_feature_slice = input_features[0:INPUT_EXAMPLE_SAMPLE_ROWS]
+    return input_feature_slice
+
+
+def _extract_sample_numpy_dict(
+    input_numpy_features_dict: dict[str, np.ndarray],
+) -> Union[dict[str, np.ndarray], np.ndarray]:
+    """
+    Extracts `INPUT_EXAMPLE_SAMPLE_ROWS` sample from next_input
+    as numpy array of dict(str -> ndarray) type.
+
+    Args:
+        input_numpy_features_dict: A tensor or numpy array
+
+    Returns:
+        A slice (limit `INPUT_EXAMPLE_SAMPLE_ROWS`)  of the input of same type as next_input.
+        Returns `None` if the type of `input_numpy_features_dict` is unsupported.
+
+    Examples
+    --------
+    when next_input is dict:
+    >>> input_data = {"a": np.array([1, 2, 3, 4, 5, 6, 7, 8])}
+    >>> _extract_sample_numpy_dict(input_data)
+    {'a': array([1, 2, 3, 4, 5])}
+
+    """
+    sliced_data_as_numpy = None
+    if isinstance(input_numpy_features_dict, dict):
+        sliced_data_as_numpy = {
+            k: _extract_input_example_from_tensor_or_ndarray(v)
+            for k, v in input_numpy_features_dict.items()
+        }
+    return sliced_data_as_numpy
+
+
+def _extract_input_example_from_batched_tf_dataset(
+    dataset: tensorflow.data.Dataset,
+) -> Union[np.ndarray, dict[str, np.ndarray]]:
+    """
+    Extracts sample feature tensors from the input dataset as numpy array.
+    Input Dataset's tensors must contain tuple of (features, labels) that are
+    used for tensorflow/keras train or fit methods
+
+
+    Args:
+        dataset: a tensorflow batched/unbatched dataset representing tuple of (features, labels)
+
+    Returns:
+        a numpy array of length `INPUT_EXAMPLE_SAMPLE_ROWS`
+        Returns `None` if the type of `dataset` slices are unsupported.
+
+    Examples
+    --------
+    >>> input_dataset = tensorflow.data.Dataset.from_tensor_slices(
+    ...     (
+    ...         {
+    ...             "SepalLength": np.array(list(range(0, 20))),
+    ...             "SepalWidth": np.array(list(range(0, 20))),
+    ...             "PetalLength": np.array(list(range(0, 20))),
+    ...             "PetalWidth": np.array(list(range(0, 20))),
+    ...         },
+    ...         np.array(list(range(0, 20))),
+    ...     )
+    ... ).batch(10)
+    >>> _extract_input_example_from_batched_tf_dataset(input_dataset)
+    {'SepalLength': array([0, 1, 2, 3, 4]),
+    'SepalWidth': array([0, 1, 2, 3, 4]),
+    'PetalLength': array([0, 1, 2, 3, 4]),
+    'PetalWidth': array([0, 1, 2, 3, 4])}
+
+    """
+    limited_df_iter = list(dataset.take(INPUT_EXAMPLE_SAMPLE_ROWS))
+    first_batch = limited_df_iter[0]
+    input_example_slice = None
+    if isinstance(first_batch, tuple):
+        features = first_batch[0]
+        if isinstance(features, dict):
+            input_example_slice = _extract_sample_numpy_dict(features)
+        elif isinstance(features, (np.ndarray, tensorflow.Tensor)):
+            input_example_slice = _extract_input_example_from_tensor_or_ndarray(features)
+    return input_example_slice
+
+
+def extract_input_example_from_tf_input_fn(input_fn):
+    """
+    Extracts sample data from dict (str -> ndarray),
+    ``tensorflow.Tensor`` or ``tensorflow.data.Dataset`` type.
+
+    Args:
+        input_fn: Tensorflow's input function used for train method
+
+    Returns:
+        A slice (of limit ``mlflow.utils.autologging_utils.INPUT_EXAMPLE_SAMPLE_ROWS``)
+        of the input of type `np.ndarray`.
+        Returns `None` if the return type of ``input_fn`` is unsupported.
+    """
+
+    input_training_data = input_fn()
+    input_features = None
+    if isinstance(input_training_data, tuple):
+        features = input_training_data[0]
+        if isinstance(features, dict):
+            input_features = _extract_sample_numpy_dict(features)
+        elif isinstance(features, (np.ndarray, tensorflow.Tensor)):
+            input_features = _extract_input_example_from_tensor_or_ndarray(features)
+    elif isinstance(input_training_data, tensorflow.data.Dataset):
+        input_features = _extract_input_example_from_batched_tf_dataset(input_training_data)
+    return input_features
+
+
+def extract_tf_keras_input_example(input_training_data):
+    """
+    Generates a sample ndarray or dict (str -> ndarray)
+    from the input type 'x' for keras ``fit`` or ``fit_generator``
+
+    Args:
+        input_training_data: Keras input function used for ``fit`` or ``fit_generator`` methods.
+
+    Returns:
+        a slice of type ndarray or
+        dict (str -> ndarray) limited to
+        ``mlflow.utils.autologging_utils.INPUT_EXAMPLE_SAMPLE_ROWS``.
+        Throws ``MlflowException`` exception, if input_training_data is unsupported.
+        Returns `None` if the type of input_training_data is unsupported.
+    """
+    input_data_slice = None
+    if isinstance(input_training_data, tensorflow.keras.utils.Sequence):
+        input_training_data = input_training_data[:][0]
+
+    if isinstance(input_training_data, (np.ndarray, tensorflow.Tensor)):
+        input_data_slice = _extract_input_example_from_tensor_or_ndarray(input_training_data)
+    elif isinstance(input_training_data, dict):
+        input_data_slice = _extract_sample_numpy_dict(input_training_data)
+    elif isinstance(input_training_data, tensorflow.data.Dataset):
+        input_data_slice = _extract_input_example_from_batched_tf_dataset(input_training_data)
+    return input_data_slice

venv/lib/python3.9/site-packages/mlflow/tensorflow/callback.py

@@ -0,0 +1,233 @@
+from tensorflow import keras
+from tensorflow.keras.callbacks import Callback
+
+from mlflow import log_metrics, log_params, log_text
+from mlflow.utils.autologging_utils import ExceptionSafeClass
+from mlflow.utils.checkpoint_utils import MlflowModelCheckpointCallbackBase
+
+
+class MlflowCallback(keras.callbacks.Callback, metaclass=ExceptionSafeClass):
+    """Callback for logging Tensorflow training metrics to MLflow.
+
+    This callback logs model information at training start, and logs training metrics every epoch or
+    every n steps (defined by the user) to MLflow.
+
+    Args:
+        log_every_epoch: bool, If True, log metrics every epoch. If False, log metrics every n
+            steps.
+        log_every_n_steps: int, log metrics every n steps. If None, log metrics every epoch.
+            Must be `None` if `log_every_epoch=True`.
+
+    .. code-block:: python
+        :caption: Example
+
+        from tensorflow import keras
+        import mlflow
+        import numpy as np
+
+        # Prepare data for a 2-class classification.
+        data = tf.random.uniform([8, 28, 28, 3])
+        label = tf.convert_to_tensor(np.random.randint(2, size=8))
+
+        model = keras.Sequential(
+            [
+                keras.Input([28, 28, 3]),
+                keras.layers.Flatten(),
+                keras.layers.Dense(2),
+            ]
+        )
+
+        model.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(0.001),
+            metrics=[keras.metrics.SparseCategoricalAccuracy()],
+        )
+
+        with mlflow.start_run() as run:
+            model.fit(
+                data,
+                label,
+                batch_size=4,
+                epochs=2,
+                callbacks=[mlflow.keras.MlflowCallback(run)],
+            )
+    """
+
+    def __init__(self, log_every_epoch=True, log_every_n_steps=None):
+        self.log_every_epoch = log_every_epoch
+        self.log_every_n_steps = log_every_n_steps
+
+        if log_every_epoch and log_every_n_steps is not None:
+            raise ValueError(
+                "`log_every_n_steps` must be None if `log_every_epoch=True`, received "
+                f"`log_every_epoch={log_every_epoch}` and `log_every_n_steps={log_every_n_steps}`."
+            )
+        if not log_every_epoch and log_every_n_steps is None:
+            raise ValueError(
+                "`log_every_n_steps` must be specified if `log_every_epoch=False`, received"
+                "`log_every_n_steps=False` and `log_every_n_steps=None`."
+            )
+
+    def on_train_begin(self, logs=None):
+        """Log model architecture and optimizer configuration when training begins."""
+        config = self.model.optimizer.get_config()
+        log_params({f"opt_{k}": v for k, v in config.items()})
+
+        model_summary = []
+
+        def print_fn(line, *args, **kwargs):
+            model_summary.append(line)
+
+        self.model.summary(print_fn=print_fn)
+        summary = "\n".join(model_summary)
+        log_text(summary, artifact_file="model_summary.txt")
+
+    def on_epoch_end(self, epoch, logs=None):
+        """Log metrics at the end of each epoch."""
+        if not self.log_every_epoch or logs is None:
+            return
+        log_metrics(logs, step=epoch, synchronous=False)
+
+    def on_batch_end(self, batch, logs=None):
+        """Log metrics at the end of each batch with user specified frequency."""
+        if self.log_every_n_steps is None or logs is None:
+            return
+        current_iteration = int(self.model.optimizer.iterations.numpy())
+
+        if current_iteration % self.log_every_n_steps == 0:
+            log_metrics(logs, step=current_iteration, synchronous=False)
+
+    def on_test_end(self, logs=None):
+        """Log validation metrics at validation end."""
+        if logs is None:
+            return
+        metrics = {"validation_" + k: v for k, v in logs.items()}
+        log_metrics(metrics, synchronous=False)
+
+
+class MlflowModelCheckpointCallback(Callback, MlflowModelCheckpointCallbackBase):
+    """Callback for automatic Keras model checkpointing to MLflow.
+
+    Args:
+        monitor: In automatic model checkpointing, the metric name to monitor if
+            you set `model_checkpoint_save_best_only` to True.
+        save_best_only: If True, automatic model checkpointing only saves when
+            the model is considered the "best" model according to the quantity
+            monitored and previous checkpoint model is overwritten.
+        mode: one of {"min", "max"}. In automatic model checkpointing,
+            if save_best_only=True, the decision to overwrite the current save file is made
+            based on either the maximization or the minimization of the monitored quantity.
+        save_weights_only: In automatic model checkpointing, if True, then
+            only the model’s weights will be saved. Otherwise, the optimizer states,
+            lr-scheduler states, etc are added in the checkpoint too.
+        save_freq: `"epoch"` or integer. When using `"epoch"`, the callback
+            saves the model after each epoch. When using integer, the callback
+            saves the model at end of this many batches. Note that if the saving isn't
+            aligned to epochs, the monitored metric may potentially be less reliable (it
+            could reflect as little as 1 batch, since the metrics get reset
+            every epoch). Defaults to `"epoch"`.
+
+    .. code-block:: python
+        :caption: Example
+
+        from tensorflow import keras
+        import tensorflow as tf
+        import mlflow
+        import numpy as np
+        from mlflow.tensorflow import MlflowModelCheckpointCallback
+
+        # Prepare data for a 2-class classification.
+        data = tf.random.uniform([8, 28, 28, 3])
+        label = tf.convert_to_tensor(np.random.randint(2, size=8))
+
+        model = keras.Sequential(
+            [
+                keras.Input([28, 28, 3]),
+                keras.layers.Flatten(),
+                keras.layers.Dense(2),
+            ]
+        )
+
+        model.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(0.001),
+            metrics=[keras.metrics.SparseCategoricalAccuracy()],
+        )
+
+        mlflow_checkpoint_callback = MlflowModelCheckpointCallback(
+            monitor="sparse_categorical_accuracy",
+            mode="max",
+            save_best_only=True,
+            save_weights_only=False,
+            save_freq="epoch",
+        )
+
+        with mlflow.start_run() as run:
+            model.fit(
+                data,
+                label,
+                batch_size=4,
+                epochs=2,
+                callbacks=[mlflow_checkpoint_callback],
+            )
+    """
+
+    def __init__(
+        self,
+        monitor="val_loss",
+        mode="min",
+        save_best_only=True,
+        save_weights_only=False,
+        save_freq="epoch",
+    ):
+        Callback.__init__(self)
+        MlflowModelCheckpointCallbackBase.__init__(
+            self,
+            checkpoint_file_suffix=".h5",
+            monitor=monitor,
+            mode=mode,
+            save_best_only=save_best_only,
+            save_weights_only=save_weights_only,
+            save_freq=save_freq,
+        )
+        self.trainer = None
+        self.current_epoch = None
+        self._last_batch_seen = 0
+        self.global_step = 0
+        self.global_step_last_saving = 0
+
+    def save_checkpoint(self, filepath: str):
+        if self.save_weights_only:
+            self.model.save_weights(filepath, overwrite=True)
+        else:
+            self.model.save(filepath, overwrite=True)
+
+    def on_epoch_begin(self, epoch, logs=None):
+        self.current_epoch = epoch
+
+    def on_train_batch_end(self, batch, logs=None):
+        # Note that `on_train_batch_end` might be invoked by every N train steps,
+        # (controlled by `steps_per_execution` argument in `model.compile` method).
+        # the following logic is similar to
+        # https://github.com/keras-team/keras/blob/e6e62405fa1b4444102601636d871610d91e5783/keras/callbacks/model_checkpoint.py#L212
+        add_batches = batch + 1 if batch <= self._last_batch_seen else batch - self._last_batch_seen
+        self._last_batch_seen = batch
+
+        self.global_step += add_batches
+
+        if isinstance(self.save_freq, int):
+            if self.global_step - self.global_step_last_saving >= self.save_freq:
+                self.check_and_save_checkpoint_if_needed(
+                    current_epoch=self.current_epoch,
+                    global_step=self.global_step,
+                    metric_dict={k: float(v) for k, v in logs.items()},
+                )
+                self.global_step_last_saving = self.global_step
+
+    def on_epoch_end(self, epoch, logs=None):
+        if self.save_freq == "epoch":
+            self.check_and_save_checkpoint_if_needed(
+                current_epoch=self.current_epoch,
+                global_step=self.global_step,
+                metric_dict={k: float(v) for k, v in logs.items()},
+            )