mlflow

MLflow logging utilities for PyMC models.

This module provides utilities to log various aspects of PyMC models to MLflow which is then extended to PyMC-Marketing models.

Autologging is supported for PyMC models and PyMC-Marketing models. This including logging of sampler diagnostics, model information, data used in the model, and InferenceData objects.

The autologging can be enabled by calling the autolog function. This function patches the pymc.sample and MMM.fit calls to log the required information.

Examples

Autologging for a PyMC model:

import mlflow

import pymc as pm

import pymc_marketing.mlflow

pymc_marketing.mlflow.autolog()

# Usual PyMC model code
with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=1, observed=[1, 2, 3])

# Incorporate into MLflow workflow
mlflow.set_experiment("PyMC Experiment")

with mlflow.start_run():
    idata = pm.sample(model=model)

Autologging for a PyMC-Marketing model:

import pandas as pd

import mlflow

from pymc_marketing.mmm import (
    GeometricAdstock,
    LogisticSaturation,
    MMM,
)
import pymc_marketing.mlflow

pymc_marketing.mlflow.autolog(log_mmm=True)

# Usual PyMC-Marketing model code

data_url = "https://raw.githubusercontent.com/pymc-labs/pymc-marketing/main/data/mmm_example.csv"
data = pd.read_csv(data_url, parse_dates=["date_week"])

X = data.drop("y",axis=1)
y = data["y"]

mmm = MMM(
    adstock=GeometricAdstock(l_max=8),
    saturation=LogisticSaturation(),
    date_column="date_week",
    channel_columns=["x1", "x2"],
    control_columns=[
        "event_1",
        "event_2",
        "t",
    ],
    yearly_seasonality=2,
)

# Incorporate into MLflow workflow

mlflow.set_experiment("MMM Experiment")

with mlflow.start_run():
    idata = mmm.fit(X, y)

    # Additional specific logging
    fig = mmm.plot_components_contributions()
    mlflow.log_figure(fig, "components.png")

Functions

autolog([log_sampler_info, log_datasets, ...])	Autologging support for PyMC models and PyMC-Marketing models.
log_arviz_summary(idata, path[, var_names])	Log the ArviZ summary as an artifact on MLflow.
log_data(model, idata)	Log the data used in the model to MLflow.
log_inference_data(idata[, save_file])	Log the InferenceData to MLflow.
log_likelihood_type(model)	Save the likelihood type of the model to MLflow.
log_model_derived_info(model)	Log various model derived information to MLflow.
log_model_graph(model, path)	Log the model graph PDF as artifact on MLflow.
log_sample_diagnostics(idata[, tune])	Log sample diagnostics to MLflow.
log_types_of_parameters(model)	Log the types of parameters in a PyMC model to MLflow.