MLflow experiment tracking exporter
Stormlog now exports memory profiles and profiling metadata to MLflow, enabling integrated tracking of GPU memory alongside model metrics, parameters, and artifacts in your MLflow experiments.
- Log memory traces and summary statistics directly to MLflow runs
- Automatic model registry integration for memory profiles
- Feature parity with Weights & Biases exporter for experiment tracking
- Tag and filter runs by memory profiles in MLflow UI
python
from stormlog.exporters import MLflowExporter
from stormlog.pytorch import GPUMemoryProfiler
import mlflow
mlflow.start_run()
profiler = GPUMemoryProfiler()
exporter = MLflowExporter()
with profiler.profile_context("training"):
# Your training loop
for batch in dataloader:
model(batch).backward()
exporter.export(profiler.trace(), "gpu_memory_profile")
mlflow.end_run()
