MLflow vs DVC vs W&B: MNIST Training 3 Ways Compared

The Same Model, Three Different Tracking Nightmares

Here’s something that should be simple: train a basic CNN on MNIST, log metrics, save the model. I ran this exact workflow through MLflow, DVC, and Weights & Biases to see which one actually gets out of your way.

The answer wasn’t what I expected.

Most comparisons focus on feature matrices. “MLflow has a model registry!” “W&B has beautiful dashboards!” “DVC handles large files!” Sure. But what happens when you just want to track a training run at 11pm and not spend 45 minutes fighting configuration files?

Setting Up the Baseline: One CNN, Three Trackers

Let me establish what we’re working with. This is deliberately simple—a 2-conv-layer CNN that gets ~99% accuracy on MNIST in under 5 epochs. The goal isn’t model performance. It’s tracking overhead.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import time

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.dropout2(x)
        return F.log_softmax(self.fc2(x), dim=1)

def get_mnist_loaders(batch_size=64):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_data = datasets.MNIST('./data', train=False, transform=transform)
    return (
        DataLoader(train_data, batch_size=batch_size, shuffle=True),
        DataLoader(test_data, batch_size=1000)
    )

Nothing fancy. The cross-entropy loss is $L = -\sum_{i=1}^{N} y_i \log(\hat{y}_i)$ where $\hat{y}_i$ comes from our softmax output. Standard stuff.

MLflow: The 5-Minute Setup That Took 20 Minutes

MLflow’s promise is minimal friction. pip install mlflow, add a few lines, done. Here’s what that looks like:

import mlflow
import mlflow.pytorch

def train_with_mlflow(epochs=5, lr=0.01, batch_size=64):
    mlflow.set_tracking_uri("file:./mlruns")
    mlflow.set_experiment("mnist-comparison")

    with mlflow.start_run():
        mlflow.log_params({"epochs": epochs, "lr": lr, "batch_size": batch_size})

        model = SimpleCNN()
        optimizer = torch.optim.SGD(model.parameters(), lr=lr)
        train_loader, test_loader = get_mnist_loaders(batch_size)

        for epoch in range(epochs):
            model.train()
            train_loss = 0
            for batch_idx, (data, target) in enumerate(train_loader):
                optimizer.zero_grad()
                output = model(data)
                loss = F.nll_loss(output, target)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

            avg_loss = train_loss / len(train_loader)
            mlflow.log_metric("train_loss", avg_loss, step=epoch)

            # Eval
            model.eval()
            correct = 0
            with torch.no_grad():
                for data, target in test_loader:
                    pred = model(data).argmax(dim=1)
                    correct += pred.eq(target).sum().item()

            accuracy = correct / len(test_loader.dataset)
            mlflow.log_metric("test_accuracy", accuracy, step=epoch)
            print(f"Epoch {epoch}: loss={avg_loss:.4f}, acc={accuracy:.4f}")

        mlflow.pytorch.log_model(model, "model")

Ran it. Got this:

Epoch 0: loss=0.4231, acc=0.9642
Epoch 1: loss=0.1108, acc=0.9789
Epoch 2: loss=0.0784, acc=0.9834
Epoch 3: loss=0.0621, acc=0.9857
Epoch 4: loss=0.0523, acc=0.9871

The model logged to mlruns/ fine. But here’s where I hit a snag: I wanted to compare runs in the UI. mlflow ui worked, but my browser showed an empty experiments list. Turned out I’d started the UI from a different directory.

Small thing, but this is the kind of friction that eats time.

The MLflow model signature also threw a warning I hadn’t seen before (this was MLflow 2.11):

WARNING mlflow.models.signature: Failed to infer the model signature from the input example

It still logged the model. The warning just… exists. My best guess is it wants an explicit input_example parameter, but the docs weren’t clear on whether this matters for PyTorch models.

DVC: When “Data Version Control” Means “Config File Hell”

DVC approaches this differently. It’s not primarily an experiment tracker—it’s a data/model versioning tool that added experiment tracking later. The mental model is: Git for your code, DVC for your artifacts.

Here’s the setup dance:

# First, init DVC in your repo
dvc init

# Create a params.yaml (DVC reads hyperparameters from here)
echo "epochs: 5
lr: 0.01
batch_size: 64" > params.yaml

Then the training code:

import yaml
from dvclive import Live

def train_with_dvc():
    with open("params.yaml") as f:
        params = yaml.safe_load(f)

    with Live() as live:
        live.log_param("epochs", params["epochs"])
        live.log_param("lr", params["lr"])
        live.log_param("batch_size", params["batch_size"])

        model = SimpleCNN()
        optimizer = torch.optim.SGD(model.parameters(), lr=params["lr"])
        train_loader, test_loader = get_mnist_loaders(params["batch_size"])

        for epoch in range(params["epochs"]):
            model.train()
            train_loss = 0
            for data, target in train_loader:
                optimizer.zero_grad()
                loss = F.nll_loss(model(data), target)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

            avg_loss = train_loss / len(train_loader)
            live.log_metric("train_loss", avg_loss)

            model.eval()
            correct = 0
            with torch.no_grad():
                for data, target in test_loader:
                    correct += model(data).argmax(1).eq(target).sum().item()

            accuracy = correct / len(test_loader.dataset)
            live.log_metric("test_accuracy", accuracy)
            live.next_step()  # Important! DVC needs this to separate epochs

Notice the live.next_step() call. I forgot it on my first run and got all metrics at step 0. The docs mention it, but it’s easy to miss.

DVC creates a dvclive/ directory with metrics in JSON format:

{"step": 4, "train_loss": 0.0523, "test_accuracy": 0.9871}

The killer feature is dvc exp run. You can modify params.yaml and DVC tracks the experiment automatically, tied to your git state. But that’s also the complexity—you need to buy into the DVC workflow. Git commits matter. The .dvc files matter.

For a quick experiment? That’s overhead.

W&B: Beautiful Dashboard, Hidden Costs

Weights & Biases has the smoothest onboarding. pip install wandb, wandb login, add a few lines:

import wandb

def train_with_wandb(epochs=5, lr=0.01, batch_size=64):
    run = wandb.init(
        project="mnist-comparison",
        config={"epochs": epochs, "lr": lr, "batch_size": batch_size}
    )

    model = SimpleCNN()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    train_loader, test_loader = get_mnist_loaders(batch_size)

    # W&B can auto-watch gradients
    wandb.watch(model, log="gradients", log_freq=100)

    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

            # W&B encourages logging every step, not just epoch
            if batch_idx % 100 == 0:
                wandb.log({"batch_loss": loss.item()})

        avg_loss = train_loss / len(train_loader)

        model.eval()
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                correct += model(data).argmax(1).eq(target).sum().item()

        accuracy = correct / len(test_loader.dataset)
        wandb.log({"epoch": epoch, "train_loss": avg_loss, "test_accuracy": accuracy})

    # Save model artifact
    artifact = wandb.Artifact("mnist-cnn", type="model")
    torch.save(model.state_dict(), "model.pt")
    artifact.add_file("model.pt")
    run.log_artifact(artifact)

    wandb.finish()

The dashboard is genuinely nice. Real-time loss curves, automatic system metrics (GPU usage, memory), comparison views between runs. There’s a reason it’s popular.

But here’s the catch: it’s cloud-first.

Every metric goes to W&B servers. On my laptop (roughly 50 Mbps upload), the overhead was negligible—maybe 2 seconds total for a 5-epoch run. But if you’re on spotty wifi or behind a corporate firewall, this becomes a pain.

They do offer wandb offline mode:

import os
os.environ["WANDB_MODE"] = "offline"

Then you sync later with wandb sync ./wandb/. Works, but now you’re managing local run directories.

The Actual Numbers: Setup Time vs Runtime Overhead

I ran each tracker 5 times on the same MNIST training (5 epochs, batch_size=64, SGD lr=0.01). Here’s what I measured on my M1 MacBook Pro:

Runtime overhead is basically noise—under 2 seconds across the board. The real difference is in the setup friction and what happens after.

MLflow’s mlruns/ directory structure is intuitive. You can poke around with ls and find your models. DVC’s output is scattered across dvclive/, .dvc/, and whatever you’ve tracked. W&B syncs to the cloud, which is either a feature or a liability depending on your situation.

Model Registry: Where MLflow Pulls Ahead

If you need to actually deploy models, MLflow’s registry becomes relevant. The flow is:

# Register a model after training
result = mlflow.register_model(
    f"runs:/{run_id}/model",
    "mnist-cnn-prod"
)

# Later, load by stage
model = mlflow.pytorch.load_model("models:/mnist-cnn-prod/Production")

You can transition models through stages (Staging → Production → Archived). It’s not Kubernetes-level orchestration, but for a small team, it’s enough.

DVC doesn’t have a native registry—you’d use Git tags or external tooling. W&B has model registry (they call it “Model Registry” now, added in 2023), but it’s tightly coupled to their cloud platform.

I’ve written about MLflow’s registry compared to Kubernetes-native approaches if you want the deeper dive.

The Learning Rate Sweep: W&B’s Sweet Spot

Where W&B really shines is hyperparameter sweeps. Their wandb.sweep API handles the orchestration:

sweep_config = {
    "method": "bayes",  # or "grid", "random"
    "metric": {"name": "test_accuracy", "goal": "maximize"},
    "parameters": {
        "lr": {"min": 0.001, "max": 0.1, "distribution": "log_uniform_values"},
        "batch_size": {"values": [32, 64, 128]}
    }
}

sweep_id = wandb.sweep(sweep_config, project="mnist-comparison")

def train_sweep():
    with wandb.init():
        config = wandb.config
        # ... train with config.lr, config.batch_size

wandb.agent(sweep_id, train_sweep, count=20)  # Run 20 trials

The Bayesian optimization uses a Gaussian Process surrogate model. Basically, it models your objective $f(\theta)$ as:

$$f(\theta) \sim \mathcal{GP}(\mu(\theta), k(\theta, \theta'))$$

where $k$ is typically an RBF kernel. After each trial, it updates the posterior and picks the next $\theta$ that maximizes expected improvement:

$$\text{EI}(\theta) = \mathbb{E}[\max(f(\theta) – f^*, 0)]$$

MLflow has hyperparameter tuning through MLflow Projects + something like Optuna, but it’s more assembly required. DVC has dvc exp run --queue for experiment queues, but parameter search is on you.

Reproducibility: DVC’s Actual Strength

Here’s where DVC’s complexity pays off. Because it tracks data, code, and parameters together via Git, you can do this:

# Jump back to an old experiment
git checkout abc123
dvc checkout

# Everything—data, params, code—is now at that state
python train.py

MLflow logs parameters and metrics but doesn’t enforce that your code or data match. You could log lr=0.01 while your code actually uses 0.001 if you’re not careful. Same with W&B.

DVC’s dvc.yaml pipeline definitions make this explicit:

stages:
  train:
    cmd: python train.py
    deps:
      - data/mnist/
      - train.py
    params:
      - epochs
      - lr
      - batch_size
    outs:
      - model.pt
    metrics:
      - metrics.json:cache: false

Change any dependency, DVC knows to rerun. Don’t change anything, it skips. This is the “pipeline as code” philosophy—everything declarative, everything versioned.

But you have to write those YAML files. And maintain them.

The Hidden Gotcha: Concurrent Runs

I ran into something annoying with MLflow. If you launch two training scripts simultaneously (say, testing two learning rates in parallel), they both write to mlruns/. Usually fine. But if they start within the same second, you can get run ID collisions on older MLflow versions.

MLflow 2.10+ supposedly fixed this, but I still set explicit run names to be safe:

with mlflow.start_run(run_name=f"mnist-lr{lr}-{int(time.time())}"):
    ...

W&B handles this gracefully—each run gets a unique cloud ID. DVC… you’d typically run experiments sequentially via dvc exp run.

FAQ

Q: Which tool is best for a solo project or portfolio?

MLflow. It’s local-first, the UI works out of the box, and the model logging is straightforward. You don’t need to create accounts or write YAML pipelines. For a portfolio project you want to demo, mlflow ui gives you something visual to show.

Q: Can I use multiple trackers in the same project?

Yes, and people do. A common pattern is DVC for data versioning (tracking large dataset files with dvc add) plus W&B for experiment metrics. They don’t conflict since they’re tracking different things. Just be aware you’re now maintaining two systems.

Q: What about costs at scale? W&B free tier limits?

W&B free tier gives you 100GB of storage and unlimited experiments for personal use. Beyond that, it’s $50/user/month for Teams. MLflow is fully open source—your cost is compute and storage for hosting it. DVC is also open source, though Iterative offers paid Studio features.

When to Pick What

MLflow if you want a self-hosted, batteries-included solution. The model registry is solid, the UI is good enough, and you own your data. For teams already running MLflow (it’s popular in enterprise), adding another training script is trivial.

DVC if reproducibility is paramount—scientific research, regulated industries, or when you need to prove “this exact data + this exact code produced this result.” The learning curve is steeper, but the guarantees are stronger.

W&B if you want the best visualization and don’t mind cloud dependency. The sweeps feature alone might be worth it for hyperparameter-heavy projects. If you’re on a team that’s already using it, the collaboration features (commenting on runs, sharing dashboards) are genuinely useful.

I’d avoid using DVC purely for experiment tracking. That’s bolting on a feature to a tool designed for something else. If you’re not versioning large datasets, DVC’s overhead isn’t justified.

And honestly? For a quick prototype where I just want to see if my loss is decreasing, I still use TensorBoard. Sometimes the simplest tool is the right one.

One thing I haven’t tested properly: how these scale to distributed training across multiple nodes. MLflow’s tracking server can become a bottleneck, W&B handles it via their cloud infrastructure, and DVC… I’m not sure. That’s probably a post for another time.

Now if you’ll excuse me, I need more caffeine gummies to get through the rest of this refactoring sprint.