# How to Implement Cross-Framework Model Logging (PyTorch & scikit-learn) with CMF

> Implement cross-framework model logging for PyTorch and scikit-learn using CMF. Unified metadata tracking for ML models without framework-specific logic. Learn more!

- Repository: [Hewlett Packard Enterprise/cmf](https://github.com/hewlettpackard/cmf)
- Tags: how-to-guide
- Published: 2026-03-03

---

**CMF provides a framework-agnostic auto-logging system that unifies PyTorch, scikit-learn, and other ML libraries by treating model files as URIs wrapped in the `MLModel` class, enabling unified metadata tracking without framework-specific serialization logic.**

Managing machine learning metadata across different frameworks often requires fragmented logging code. The CMF (Cross-Model Framework) library from Hewlett Packard Enterprise solves this through a unified auto-logging interface that treats all model formats identically. This guide demonstrates how to implement cross-framework model logging using CMF's decorator-based API to track both PyTorch and scikit-learn artifacts in the same metadata store.

## The Framework-Agnostic Architecture

### Core Components in auto_logging_v01.py

The auto-logging system in [`cmflib/contrib/auto_logging_v01.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/contrib/auto_logging_v01.py) provides the abstraction layer that makes cross-framework compatibility possible:

- **Artifact** – Base class representing any URI-addressable resource with optional parameters
- **MLModel** – Subclass that wraps model file paths without inspecting contents  
- **@step** – Decorator that intercepts function returns and automatically invokes `Cmf.log_model()`
- **Cmf** – Low-level API in [`cmflib/cmf.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/cmf.py) handling MLMD (ML Metadata) storage and optional Neo4j graph integration

The architecture relies on **file-based abstraction**: `MLModel` stores only the file URI and optional metadata, never parsing the serialized contents. Whether you save a scikit-learn estimator with `pickle.dump()` or a PyTorch state dict with `torch.save()`, CMF logs both as identical artifact types.

### Execution Flow

When you decorate a function with `@step()`, the following occurs:

1. The decorator creates a `Cmf` instance that initializes the execution context
2. Your function receives `Context`, `Parameters`, and input `Artifact` objects
3. You persist models using native framework APIs (pickle, torch.save, etc.)
4. You return a dictionary mapping names to `MLModel` or other `Artifact` objects
5. The decorator automatically calls `Cmf.log_model()` for every returned artifact before the function exits

## Implementing Cross-Framework Model Logging

### scikit-learn Example (DecisionTree)

The following pipeline step from [`examples/auto_logging_v01/pipeline/train_sklearn.py`](https://github.com/hewlettpackard/cmf/blob/main/examples/auto_logging_v01/pipeline/train_sklearn.py) demonstrates logging a DecisionTree classifier:

```python

# examples/auto_logging_v01/pipeline/train_sklearn.py

import pickle, typing as t
from pathlib import Path
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score

from cmflib.cmf import Cmf
from cmflib.contrib.auto_logging_v01 import (
    Context, Dataset, ExecutionMetrics, MLModel, Parameters, cli_run, prepare_workspace, step,
)

@step()
def train_sklearn(
    ctx: Context, params: Parameters, train_dataset: Dataset
) -> t.Dict[str, t.Union[MLModel, ExecutionMetrics]]:
    """Train a scikit-learn model and log it."""
    # Load the dataset (pickled dict with keys "x", "y")

    with open(train_dataset.uri, "rb") as f:
        data = pickle.load(f)

    # Fit the model

    clf = DecisionTreeClassifier()
    clf.fit(data["x"], data["y"])

    # Compute a metric we want to record

    acc = accuracy_score(data["y"], clf.predict(data["y"]))

    # Persist the model – any format works; we use pickle here

    workspace = prepare_workspace(ctx)
    model_path = workspace / "dt_model.pkl"
    with open(model_path, "wb") as f:
        pickle.dump(clf, f)

    # Return artifacts. CMF will log the model and the metric automatically.

    cmf: Cmf = ctx["cmf"]
    return {
        "model": MLModel(model_path),
        "metrics": ExecutionMetrics(
            uri=str(cmf.execution.id) + "/metrics/train", name="train", params={"accuracy": acc}
        ),
    }

if __name__ == "__main__":
    # Example launch:

    #   python train_sklearn.py train_dataset=workspace/train.pkl

    cli_run(train_sklearn)

```

**Key implementation details:**

- Uses native `pickle.dump()` for serialization, requiring no CMF-specific export format
- Returns `MLModel(model_path)` to trigger automatic logging via the decorator
- The `@step` decorator handles `Cmf.log_model()` invocation automatically without boilerplate

### PyTorch Example (ResNet-18)

This example from [`examples/trustworthyAI-cmf/src/train_torch.py`](https://github.com/hewlettpackard/cmf/blob/main/examples/trustworthyAI-cmf/src/train_torch.py) logs a ResNet-18 model using identical CMF patterns:

```python

# examples/trustworthyAI-cmf/src/train_torch.py

import torch, typing as t
from pathlib import Path
from torchvision import models

from cmflib.cmf import Cmf
from cmflib.contrib.auto_logging_v01 import (
    Context, Dataset, ExecutionMetrics, MLModel, Parameters, cli_run, prepare_workspace, step,
)

@step()
def train_torch(
    ctx: Context, params: Parameters, train_dataset: Dataset
) -> t.Dict[str, t.Union[MLModel, ExecutionMetrics]]:
    """Train a simple ResNet-18 with PyTorch and log the artefacts."""
    # Load pre-processed tensors (saved as .pt by the example)

    train_data = torch.load(train_dataset.uri)   # expects a dict with "inputs", "labels"

    # Build the model

    model = models.resnet18(pretrained=False)
    model.fc = torch.nn.Linear(model.fc.in_features, train_data["labels"].max().item() + 1)

    # Training loop (very short for demo)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()

    model.train()
    for epoch in range(2):                     # tiny demo epoch count

        optimizer.zero_grad()
        inputs = train_data["inputs"].to(device)
        labels = train_data["labels"].to(device)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # Record a metric

    acc = (outputs.argmax(dim=1) == labels).float().mean().item()

    # Persist the model using torch.save (any extension works)

    workspace = prepare_workspace(ctx)
    model_path = workspace / "resnet18.pt"
    torch.save(model.state_dict(), model_path)

    # Return artifacts – CMF will invoke `Cmf.log_model` under the hood.

    cmf: Cmf = ctx["cmf"]
    return {
        "model": MLModel(model_path),
        "metrics": ExecutionMetrics(
            uri=str(cmf.execution.id) + "/metrics/train", name="train", params={"accuracy": acc}
        ),
    }

if __name__ == "__main__":
    # Example launch:

    #   python train_torch.py train_dataset=workspace/train.pt

    cli_run(train_torch)

```

**Implementation notes:**

- Native `torch.save()` persists the state dict to `.pt` format
- Same `MLModel` wrapper works despite different serialization format from scikit-learn
- No changes to CMF configuration required when switching between deep learning and classical ML frameworks

## Retrieving Logged Models Across Frameworks

Both frameworks write identical metadata schemas to the MLMD store, enabling unified retrieval using `get_latest_artifact()`:

```python
from cmflib.cmf import Cmf

cmf = Cmf(filepath="mlmd", pipeline_name="my_pipeline")

# Find the latest model for a given stage

model_art = cmf.get_latest_artifact(artifact_type="MLModel", stage="train")
if model_art.uri.endswith(".pt"):
    # PyTorch – load state dict

    import torch, torchvision
    model = torchvision.models.resnet18()
    model.load_state_dict(torch.load(model_art.uri))
else:
    # scikit-learn – use pickle

    import pickle
    model = pickle.load(open(model_art.uri, "rb"))

```

The **same retrieval logic** works for both frameworks because the artifact metadata (URI, custom properties) is stored in MLMD and can be queried via the CMF API regardless of the original training library.

## Summary

- **File-based abstraction** allows CMF to log any framework by treating models as URIs rather than specific object types
- The **MLModel** class in [`cmflib/contrib/auto_logging_v01.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/contrib/auto_logging_v01.py) requires only a file path, not framework-specific parsing logic
- The **@step** decorator automates metadata capture by intercepting return values and invoking `log_model()` automatically
- **Native serialization** (pickle, torch.save, etc.) is preserved—no forced format conversions or custom exporters required
- Unified retrieval via `get_latest_artifact()` works across PyTorch, scikit-learn, XGBoost, and future frameworks without query modifications

## Frequently Asked Questions

### Does CMF require specific model formats to enable cross-framework logging?

No. CMF's `MLModel` class only records the file URI and custom properties in the metadata store. It does not inspect or validate the file contents, allowing you to use pickle, `.pt` files, TensorFlow SavedModel directories, or any other serialization format native to your ML library.

### How does the @step decorator know which artifacts to log?

The decorator inspects the return value of your function in [`cmflib/contrib/auto_logging_v01.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/contrib/auto_logging_v01.py). When you return a dictionary containing `Artifact` subclasses (like `MLModel` or `ExecutionMetrics`), it automatically invokes the appropriate `Cmf.log_*` methods (e.g., `log_model()`, `log_execution_metrics()`) before completing the function execution.

### Can I mix PyTorch and scikit-learn steps in the same CMF pipeline?

Yes. Because both frameworks use the same metadata schema via `MLModel` and identical storage mechanisms in [`cmflib/cmf.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/cmf.py), you can chain PyTorch preprocessing steps with scikit-learn training steps (or vice versa) in a single pipeline. The lineage graph will correctly track dependencies across framework boundaries.

### Where does CMF store the metadata for logged models?

By default, CMF stores metadata in an SQLite-backed MLMD (ML Metadata) database specified by the `filepath` parameter when initializing `Cmf`. Optional Neo4j graph integration is available via [`cmflib/graph_wrapper.py`](https://github.com/hewlettpackard/cmf/blob/main/cmflib/graph_wrapper.py) for complex lineage queries and visualizations.