Applying magnitude pruning to an MNIST model

In this tutorial, we will be providing a basic introduction to pruning a model with CoreAI-Opt.

After the end of this tutorial, you should be familiar with the following:

1. How to apply CoreAI-Opt’s post-training magnitude pruning

2. How to apply CoreAI-Opt’s magnitude pruning with a sparsity schedule and fine-tuning

3. How to export CoreAI-Opt pruned models to Core AI

Table of Contents:

• Setup

  • MNIST Dataset download

  • Model definition

  • Training and Evaluation

  • Train baseline model

• Post-Training Magnitude Pruning

• Magnitude Pruning with Fine-Tuning

• Export to Core AI

Setup

We will be using a basic CNN model and train it on the MNIST dataset and observe its final accuracy.

Once we train this CNN model, we will apply magnitude pruning to it using coreai-opt, starting with a post-training pass (no fine-tuning), and then moving to a scheduled pass that ramps up sparsity while fine-tuning to recover accuracy.

import random
from pathlib import Path

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchinfo import summary
from torchvision import datasets, transforms

# Use the MPS (Apple Silicon GPU) backend when available; otherwise fall back to CPU.
if torch.backends.mps.is_available():
    DEVICE = torch.device("mps")
else:
    DEVICE = torch.device("cpu")

SEED = 1976

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

<torch._C.Generator at 0x1127a9b50>

# Used to save intermediate results and datasets
SAVE_DIRECTORY = "."

MNIST Dataset download

Helper to download the MNIST dataset with standard normalization applied.

def mnist_transforms() -> transforms.Compose:
    return transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])


def download_mnist_dataset(
    download_path: Path, transform: transforms.Compose | None = None
) -> tuple[datasets.MNIST, datasets.MNIST]:
    if transform is None:
        transform = mnist_transforms()
    train = datasets.MNIST(download_path, train=True, download=True, transform=transform)
    test = datasets.MNIST(download_path, train=False, download=True, transform=transform)
    return train, test

Model definition

A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier.

class MnistNetwork(nn.Module):
    def __init__(self, num_classes: int = 10, state_dict: dict | None = None) -> None:
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 12, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2, padding=0),
            nn.Flatten(),
            nn.Linear(2352, num_classes),
        )
        if state_dict is not None:
            self.load_state_dict(state_dict)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

Training and Evaluation

Standard PyTorch training loop and evaluation function that computes accuracy.

def train_step(model, optimizer, loss_fn, inputs, ground_truth) -> float:
    model.train()
    device = next(model.parameters()).device
    inputs = inputs.to(device)
    ground_truth = ground_truth.to(device)
    optimizer.zero_grad()
    predictions = model(inputs)
    loss = loss_fn(predictions, ground_truth)
    loss.backward()
    optimizer.step()
    return loss.item()


def train_epoch(model, train_loader, optimizer, loss_fn) -> float:
    total_loss = 0.0
    for inputs, ground_truth in train_loader:
        loss = train_step(model, optimizer, loss_fn, inputs, ground_truth)
        total_loss += loss
    return total_loss / len(train_loader)


def create_adam_optimizer(model: nn.Module, lr: float = 1e-3) -> torch.optim.Adam:
    return torch.optim.Adam(model.parameters(), lr=lr)


def eval_model(model: nn.Module, test_dataloader: DataLoader) -> float:
    model.eval()
    device = next(model.parameters()).device
    num_correct = 0
    total = 0
    with torch.no_grad():
        for inputs, ground_truth in test_dataloader:
            inputs = inputs.to(device)
            ground_truth = ground_truth.to(device)
            predictions = model(inputs)
            _, predicted = torch.max(predictions.data, 1)
            total += ground_truth.size(0)
            num_correct += (predicted == ground_truth).sum().item()
    return num_correct / total

# Download and instantiate datasets
DOWNLOAD_PATH = Path(SAVE_DIRECTORY) / ".mnist_dataset"

train_dataset, test_dataset = download_mnist_dataset(
    download_path=DOWNLOAD_PATH
)

BATCH_SIZE = 128

# Instantiate dataloaders
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE)

The CNN model used for this tutorial contains a single Conv2d, ReLU, MaxPool, Flatten, and Linear layer. Here’s the structure:

basic_cnn_model = MnistNetwork(num_classes=10)

# Print summary of model
summary(basic_cnn_model, input_size=(1, 1, 28, 28))

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
MnistNetwork                             [1, 10]                   --
├─Sequential: 1-1                        [1, 10]                   --
│    └─Conv2d: 2-1                       [1, 12, 28, 28]           120
│    └─ReLU: 2-2                         [1, 12, 28, 28]           --
│    └─MaxPool2d: 2-3                    [1, 12, 14, 14]           --
│    └─Flatten: 2-4                      [1, 2352]                 --
│    └─Linear: 2-5                       [1, 10]                   23,530
==========================================================================================
Total params: 23,650
Trainable params: 23,650
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 0.12
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.08
Params size (MB): 0.09
Estimated Total Size (MB): 0.17
==========================================================================================

Train baseline model

Let’s train this model so we can get a baseline accuracy. We save the trained weights so we can reload them for each pruning experiment.

EPOCHS = 10

loss_fn = torch.nn.CrossEntropyLoss()
optimizer = create_adam_optimizer(basic_cnn_model)

basic_cnn_model = basic_cnn_model.to(DEVICE)

epoch_results = []
for epoch in range(EPOCHS):
    epoch_avg_loss = train_epoch(
        model=basic_cnn_model, train_loader=train_loader, optimizer=optimizer, loss_fn=loss_fn
    )
    epoch_results.append(f"  Epoch {epoch + 1}: loss={epoch_avg_loss:.4f}")

basic_cnn_model = basic_cnn_model.cpu().eval()
print("\n".join(epoch_results))

# Save trained weights for reuse in pruning sections
pretrained_state_dict = basic_cnn_model.state_dict()

  Epoch 1: loss=0.3126
  Epoch 2: loss=0.1170
  Epoch 3: loss=0.0821
  Epoch 4: loss=0.0669
  Epoch 5: loss=0.0590
  Epoch 6: loss=0.0522
  Epoch 7: loss=0.0481
  Epoch 8: loss=0.0452
  Epoch 9: loss=0.0414
  Epoch 10: loss=0.0379

accuracy = eval_model(basic_cnn_model, test_loader)
print(f"Baseline accuracy: {accuracy:.4f}")

Baseline accuracy: 0.9798

Post-Training Magnitude Pruning

Magnitude pruning sparsifies a model by zeroing out the smallest-magnitude weights, up to a target_sparsity. Post-training pruning applies the full sparsity in a single shot during the prepare() call — no calibration data or fine-tuning is required.

Unless a model’s weights are already close to zero, post-training pruning will usually degrade accuracy. It’s most useful as a quick way to see the effect of sparsity before committing to a fine-tuning workflow.

from coreai_opt.pruning import (
    MagnitudePruner,
    MagnitudePrunerConfig,
    ModuleMagnitudePrunerConfig,
    PruningSpec,
)

example_inputs = (torch.randn(1, 1, 28, 28),)

For this tutorial, we’ll apply 50% unstructured magnitude pruning (individual elements, not whole channels) via PruningSpec. Refer to the Pruning Config page for all options.

post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)
post_training_model.eval()

# 50% unstructured magnitude pruning on every supported weight.
post_training_config = MagnitudePrunerConfig(
    global_config=ModuleMagnitudePrunerConfig(
        op_state_spec={"weight": PruningSpec(target_sparsity=0.5)},
    ),
)

post_training_pruner = MagnitudePruner(post_training_model, post_training_config)
post_training_prepared = post_training_pruner.prepare(example_inputs)
print("Prepared post-training magnitude pruning at 50% sparsity")

Prepared post-training magnitude pruning at 50% sparsity

After calling prepare(), 50% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately.

post_training_accuracy = eval_model(post_training_prepared, test_loader)
print(f"Post-training pruning accuracy: {post_training_accuracy:.4f}")

Post-training pruning accuracy: 0.8805

Magnitude Pruning with Fine-Tuning

In most cases, fine-tuning is required to recover accuracy after pruning. Instead of applying the full sparsity in one shot, we configure a sparsity_schedule on the module config and call pruner.step() once per epoch to gradually ramp up sparsity while the model keeps training.

Here we target 50% sparsity, ramped in via a PolynomialDecaySchedule over the 5 epochs of fine-tuning.

from coreai_opt.pruning.config import PolynomialDecaySchedule

FINE_TUNE_EPOCHS = 5

scheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)

scheduled_config = MagnitudePrunerConfig(
    global_config=ModuleMagnitudePrunerConfig(
        op_state_spec={"weight": PruningSpec(target_sparsity=0.5)},
        sparsity_schedule=PolynomialDecaySchedule(
            begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0
        ),
    ),
)

scheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)
scheduled_prepared = scheduled_pruner.prepare(example_inputs)
print(f"Prepared scheduled magnitude pruning targeting 50% sparsity over {FINE_TUNE_EPOCHS} epochs")

Prepared scheduled magnitude pruning targeting 50% sparsity over 5 epochs

We now fine-tune the model while incrementing the sparsity schedule. The pruner.step() call at the end of each epoch advances the schedule and recomputes the pruning masks against the current weight magnitudes for the next sparsity level.

fine_tune_optimizer = torch.optim.SGD(scheduled_prepared.parameters(), lr=1e-3)

fine_tune_results = []
for epoch in range(FINE_TUNE_EPOCHS):
    epoch_avg_loss = train_epoch(
        model=scheduled_prepared,
        train_loader=train_loader,
        optimizer=fine_tune_optimizer,
        loss_fn=loss_fn,
    )
    scheduled_pruner.step()
    epoch_accuracy = eval_model(scheduled_prepared, test_loader)
    fine_tune_results.append(
        f"  Epoch {epoch + 1}: loss={epoch_avg_loss:.4f}, accuracy={epoch_accuracy:.4f}"
    )

scheduled_prepared = scheduled_prepared.eval()
print("\n".join(fine_tune_results))

  Epoch 1: loss=0.0299, accuracy=0.9750
  Epoch 2: loss=0.0354, accuracy=0.9530
  Epoch 3: loss=0.0619, accuracy=0.9727
  Epoch 4: loss=0.0561, accuracy=0.9793
  Epoch 5: loss=0.0493, accuracy=0.9796

scheduled_accuracy = eval_model(scheduled_prepared, test_loader)
print(f"Scheduled magnitude pruning accuracy (50% sparsity, fine-tuned): {scheduled_accuracy:.4f}")

Scheduled magnitude pruning accuracy (50% sparsity, fine-tuned): 0.9796

Export to Core AI

Once the pruned model is ready, call finalize() to prepare the sparsified modules for deployment. Pass ExportBackend.CoreAI to finalize(backend=...) to target the .aimodel format produced by coreai-torch.

We’ll export the fine-tuned, scheduled-pruning model from the previous section.

from coreai_opt import ExportBackend

coreai_model = scheduled_pruner.finalize(backend=ExportBackend.CoreAI)

The export proceeds in three steps:

• Trace the model with torch.export.export() to obtain a graph representation.

• Apply cast_to_16_bit_precision() to cast remaining FP32 parameters to FP16 for optimal on-device performance.

• Convert the exported program to Core AI format using coreai-torch.TorchConverter.

import shutil

from coreai_opt.casting import cast_to_16_bit_precision
from coreai_torch import TorchConverter, get_decomp_table

exported_program = torch.export.export(coreai_model, example_inputs, strict=False)
exported_program = exported_program.run_decompositions(get_decomp_table())
cast_to_16_bit_precision(exported_program)

coreai_program = TorchConverter().add_exported_program(exported_program).to_coreai()
coreai_program.optimize()

output_path = Path(SAVE_DIRECTORY) / "exported_model.aimodel"
if output_path.exists():
    shutil.rmtree(output_path)
coreai_program.save_asset(output_path)
print(f"Exported: {output_path}")

coreai-torch 0.4.1: converting 1 program(s) to Core AI

Exported: exported_model.aimodel