From 22ab0d4af7e38502f75df327d1c31448ffe18cd9 Mon Sep 17 00:00:00 2001 From: Utkarsh Simha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:26:59 -0700 Subject: [PATCH 1/5] docs: add MNIST magnitude pruning tutorial Adds a tutorial notebook demonstrating post-training and fine-tuning-based magnitude pruning on an MNIST model, mirroring the existing quantization and palettization tutorials. Also adds the missing Pruning entry to the sidebar toctree. --- docs/src/examples/toy_models.md | 1 + docs/src/index.md | 1 + docs/src/tutorials/mnist_pruning.ipynb | 599 +++++++++++++++++++++++++ 3 files changed, 601 insertions(+) create mode 100644 docs/src/tutorials/mnist_pruning.ipynb diff --git a/docs/src/examples/toy_models.md b/docs/src/examples/toy_models.md index c2fb3af..eefc109 100644 --- a/docs/src/examples/toy_models.md +++ b/docs/src/examples/toy_models.md @@ -8,4 +8,5 @@ These tutorials demonstrate model compression techniques using simple toy models /tutorials/mnist_quantization /tutorials/mnist_palettization /tutorials/mnist_palettization_and_activation_quantization +/tutorials/mnist_pruning ``` diff --git a/docs/src/index.md b/docs/src/index.md index 92b8f7c..9d32854 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -29,6 +29,7 @@ examples/model_examples quantization/index palettization/index +pruning/index ``` ```{toctree} diff --git a/docs/src/tutorials/mnist_pruning.ipynb b/docs/src/tutorials/mnist_pruning.ipynb new file mode 100644 index 0000000..cac28aa --- /dev/null +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -0,0 +1,599 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Applying magnitude pruning to an MNIST model" + ], + "id": "cell-00" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In this tutorial, we will be providing a basic introduction to pruning a model with CoreAI-Opt.\n", + "\n", + "After the end of this tutorial, you should be familiar with the following:\n", + "1. [How to apply CoreAI-Opt's post-training magnitude pruning](#Post-Training-Magnitude-Pruning)\n", + "2. [How to apply CoreAI-Opt's magnitude pruning with a sparsity schedule and fine-tuning](#Magnitude-Pruning-with-Fine-Tuning)\n", + "3. [How to export CoreAI-Opt pruned models to Core AI](#Export-to-Core-AI)\n", + "\n", + "**Table of Contents:**\n", + "\n", + "- [Setup](#Setup)\n", + " - [MNIST Dataset download](#MNIST-Dataset-download)\n", + " - [Model definition](#Model-definition)\n", + " - [Training and Evaluation](#Training-and-Evaluation)\n", + " - [Train baseline model](#Train-baseline-model)\n", + "- [Post-Training Magnitude Pruning](#Post-Training-Magnitude-Pruning)\n", + "- [Magnitude Pruning with Fine-Tuning](#Magnitude-Pruning-with-Fine-Tuning)\n", + "- [Export to Core AI](#Export-to-Core-AI)" + ], + "id": "cell-01" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup" + ], + "id": "cell-02" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We will be using a basic CNN model and train it on the MNIST dataset and observe its final accuracy.\n", + "\n", + "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." + ], + "id": "cell-03" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import random\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import torch\n", + "from torch import nn\n", + "from torch.utils.data import DataLoader\n", + "from torchinfo import summary\n", + "from torchvision import datasets, transforms" + ], + "id": "cell-04" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Use the MPS (Apple Silicon GPU) backend when available; otherwise fall back to CPU.\n", + "if torch.backends.mps.is_available():\n", + " DEVICE = torch.device(\"mps\")\n", + "else:\n", + " DEVICE = torch.device(\"cpu\")" + ], + "id": "cell-05" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "SEED = 1976\n", + "\n", + "random.seed(SEED)\n", + "np.random.seed(SEED)\n", + "torch.manual_seed(SEED)" + ], + "id": "cell-06" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [], + "source": [ + "# Used to save intermediate results and datasets\n", + "SAVE_DIRECTORY = \".\"" + ], + "id": "cell-07" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### MNIST Dataset download", + "\n", + "Helper to download the MNIST dataset with standard normalization applied." + ], + "id": "cell-08" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def mnist_transforms() -> transforms.Compose:\n", + " return transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n", + "\n", + "\n", + "def download_mnist_dataset(\n", + " download_path: Path, transform: transforms.Compose | None = None\n", + ") -> tuple[datasets.MNIST, datasets.MNIST]:\n", + " if transform is None:\n", + " transform = mnist_transforms()\n", + " train = datasets.MNIST(download_path, train=True, download=True, transform=transform)\n", + " test = datasets.MNIST(download_path, train=False, download=True, transform=transform)\n", + " return train, test" + ], + "id": "cell-09" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Model definition", + "\n", + "A simple CNN with a single Conv2d \u2192 ReLU \u2192 MaxPool block, followed by Flatten and a Linear classifier." + ], + "id": "cell-10" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class MnistNetwork(nn.Module):\n", + " def __init__(self, num_classes: int = 10, state_dict: dict | None = None) -> None:\n", + " super().__init__()\n", + " self.model = nn.Sequential(\n", + " nn.Conv2d(1, 12, 3, padding=1),\n", + " nn.ReLU(),\n", + " nn.MaxPool2d(2, stride=2, padding=0),\n", + " nn.Flatten(),\n", + " nn.Linear(2352, num_classes),\n", + " )\n", + " if state_dict is not None:\n", + " self.load_state_dict(state_dict)\n", + "\n", + " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", + " return self.model(x)" + ], + "id": "cell-11" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Training and Evaluation", + "\n", + "Standard PyTorch training loop and evaluation function that computes accuracy." + ], + "id": "cell-12" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def train_step(model, optimizer, loss_fn, inputs, ground_truth) -> float:\n", + " model.train()\n", + " device = next(model.parameters()).device\n", + " inputs = inputs.to(device)\n", + " ground_truth = ground_truth.to(device)\n", + " optimizer.zero_grad()\n", + " predictions = model(inputs)\n", + " loss = loss_fn(predictions, ground_truth)\n", + " loss.backward()\n", + " optimizer.step()\n", + " return loss.item()\n", + "\n", + "\n", + "def train_epoch(model, train_loader, optimizer, loss_fn) -> float:\n", + " total_loss = 0.0\n", + " for inputs, ground_truth in train_loader:\n", + " loss = train_step(model, optimizer, loss_fn, inputs, ground_truth)\n", + " total_loss += loss\n", + " return total_loss / len(train_loader)\n", + "\n", + "\n", + "def create_adam_optimizer(model: nn.Module, lr: float = 1e-3) -> torch.optim.Adam:\n", + " return torch.optim.Adam(model.parameters(), lr=lr)\n", + "\n", + "\n", + "def eval_model(model: nn.Module, test_dataloader: DataLoader) -> float:\n", + " model.eval()\n", + " device = next(model.parameters()).device\n", + " num_correct = 0\n", + " total = 0\n", + " with torch.no_grad():\n", + " for inputs, ground_truth in test_dataloader:\n", + " inputs = inputs.to(device)\n", + " ground_truth = ground_truth.to(device)\n", + " predictions = model(inputs)\n", + " _, predicted = torch.max(predictions.data, 1)\n", + " total += ground_truth.size(0)\n", + " num_correct += (predicted == ground_truth).sum().item()\n", + " return num_correct / total" + ], + "id": "cell-13" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Download and instantiate datasets\n", + "DOWNLOAD_PATH = Path(SAVE_DIRECTORY) / \".mnist_dataset\"\n", + "\n", + "train_dataset, test_dataset = download_mnist_dataset(\n", + " download_path=DOWNLOAD_PATH\n", + ")" + ], + "id": "cell-14" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "BATCH_SIZE = 128\n", + "\n", + "# Instantiate dataloaders\n", + "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n", + "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE)" + ], + "id": "cell-15" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The CNN model used for this tutorial contains a single Conv2d, ReLU, MaxPool, Flatten, and Linear layer. Here's the structure:" + ], + "id": "cell-16" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "basic_cnn_model = MnistNetwork(num_classes=10)\n", + "\n", + "# Print summary of model\n", + "summary(basic_cnn_model, input_size=(1, 1, 28, 28))" + ], + "id": "cell-17" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Train baseline model" + ], + "id": "cell-18" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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." + ], + "id": "cell-19" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "EPOCHS = 10\n", + "\n", + "loss_fn = torch.nn.CrossEntropyLoss()\n", + "optimizer = create_adam_optimizer(basic_cnn_model)\n", + "\n", + "basic_cnn_model = basic_cnn_model.to(DEVICE)\n", + "\n", + "epoch_results = []\n", + "for epoch in range(EPOCHS):\n", + " epoch_avg_loss = train_epoch(\n", + " model=basic_cnn_model, train_loader=train_loader, optimizer=optimizer, loss_fn=loss_fn\n", + " )\n", + " epoch_results.append(f\" Epoch {epoch + 1}: loss={epoch_avg_loss:.4f}\")\n", + "\n", + "basic_cnn_model = basic_cnn_model.cpu().eval()\n", + "print(\"\\n\".join(epoch_results))\n", + "\n", + "# Save trained weights for reuse in pruning sections\n", + "pretrained_state_dict = basic_cnn_model.state_dict()" + ], + "id": "cell-20" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "accuracy = eval_model(basic_cnn_model, test_loader)\n", + "print(f\"Baseline accuracy: {accuracy:.4f}\")" + ], + "id": "cell-21" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Post-Training Magnitude Pruning" + ], + "id": "cell-22" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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 \u2014 no calibration data or fine-tuning is required.\n", + "\n", + "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." + ], + "id": "cell-23" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from coreai_opt.pruning import MagnitudePruner, MagnitudePrunerConfig\n", + "\n", + "example_inputs = (torch.randn(1, 1, 28, 28),)" + ], + "id": "cell-24" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For this tutorial, we'll apply the default `MagnitudePrunerConfig()`, which prunes 50% of every supported weight tensor, unstructured (individual elements, not whole channels). Refer to the [Pruning Config](../pruning/config.html) page for all options." + ], + "id": "cell-25" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", + "post_training_model.eval()\n", + "\n", + "# Default config: 50% unstructured magnitude pruning on every supported weight.\n", + "post_training_config = MagnitudePrunerConfig()\n", + "\n", + "post_training_pruner = MagnitudePruner(post_training_model, post_training_config)\n", + "post_training_prepared = post_training_pruner.prepare(example_inputs)\n", + "print(\"Prepared post-training magnitude pruning at 50% sparsity\")" + ], + "id": "cell-26" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After calling `prepare()`, 50% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately." + ], + "id": "cell-27" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "post_training_accuracy = eval_model(post_training_prepared, test_loader)\n", + "print(f\"Post-training pruning accuracy: {post_training_accuracy:.4f}\")" + ], + "id": "cell-28" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Magnitude Pruning with Fine-Tuning" + ], + "id": "cell-29" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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.\n", + "\n", + "Here we target 70% sparsity, ramped in via a `PolynomialDecaySchedule` over the 5 epochs of fine-tuning." + ], + "id": "cell-30" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from coreai_opt.pruning import ModuleMagnitudePrunerConfig, PruningSpec\n", + "from coreai_opt.pruning.config import PolynomialDecaySchedule\n", + "\n", + "FINE_TUNE_EPOCHS = 5\n", + "\n", + "scheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", + "\n", + "scheduled_config = MagnitudePrunerConfig(\n", + " global_config=ModuleMagnitudePrunerConfig(\n", + " op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n", + " sparsity_schedule=PolynomialDecaySchedule(\n", + " begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0\n", + " ),\n", + " ),\n", + ")\n", + "\n", + "scheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)\n", + "scheduled_prepared = scheduled_pruner.prepare(example_inputs)\n", + "print(f\"Prepared scheduled magnitude pruning targeting 70% sparsity over {FINE_TUNE_EPOCHS} epochs\")" + ], + "id": "cell-31" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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." + ], + "id": "cell-32" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fine_tune_optimizer = torch.optim.SGD(scheduled_prepared.parameters(), lr=1e-3)\n", + "\n", + "fine_tune_results = []\n", + "for epoch in range(FINE_TUNE_EPOCHS):\n", + " epoch_avg_loss = train_epoch(\n", + " model=scheduled_prepared,\n", + " train_loader=train_loader,\n", + " optimizer=fine_tune_optimizer,\n", + " loss_fn=loss_fn,\n", + " )\n", + " scheduled_pruner.step()\n", + " epoch_accuracy = eval_model(scheduled_prepared, test_loader)\n", + " fine_tune_results.append(\n", + " f\" Epoch {epoch + 1}: loss={epoch_avg_loss:.4f}, accuracy={epoch_accuracy:.4f}\"\n", + " )\n", + "\n", + "scheduled_prepared = scheduled_prepared.eval()\n", + "print(\"\\n\".join(fine_tune_results))" + ], + "id": "cell-33" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "scheduled_accuracy = eval_model(scheduled_prepared, test_loader)\n", + "print(f\"Scheduled magnitude pruning accuracy (70% sparsity, fine-tuned): {scheduled_accuracy:.4f}\")" + ], + "id": "cell-34" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Export to Core AI" + ], + "id": "cell-35" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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`.\n", + "\n", + "We'll export the fine-tuned, scheduled-pruning model from the previous section." + ], + "id": "cell-36" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from coreai_opt import ExportBackend\n", + "\n", + "coreai_model = scheduled_pruner.finalize(backend=ExportBackend.CoreAI)" + ], + "id": "cell-37" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The export proceeds in three steps:\n", + "\n", + "- Trace the model with `torch.export.export()` to obtain a graph representation.\n", + "- Apply `cast_to_16_bit_precision()` to cast remaining FP32 parameters to FP16 for optimal on-device performance.\n", + "- Convert the exported program to Core AI format using `coreai-torch.TorchConverter`." + ], + "id": "cell-38" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import shutil\n", + "\n", + "from coreai_opt.casting import cast_to_16_bit_precision\n", + "from coreai_torch import TorchConverter, get_decomp_table\n", + "\n", + "exported_program = torch.export.export(coreai_model, example_inputs, strict=False)\n", + "exported_program = exported_program.run_decompositions(get_decomp_table())\n", + "cast_to_16_bit_precision(exported_program)\n", + "\n", + "coreai_program = TorchConverter().add_exported_program(exported_program).to_coreai()\n", + "coreai_program.optimize()\n", + "\n", + "output_path = Path(SAVE_DIRECTORY) / \"exported_model.aimodel\"\n", + "if output_path.exists():\n", + " shutil.rmtree(output_path)\n", + "coreai_program.save_asset(output_path)\n", + "print(f\"Exported: {output_path}\")" + ], + "id": "cell-39" + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv-tutorial", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From db6ce2219a8afb57600355591c1da3f25594e05f Mon Sep 17 00:00:00 2001 From: Utkarsh Simha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:27:00 -0700 Subject: [PATCH 2/5] docs: use 70% sparsity for post-training pruning in MNIST tutorial Aligns the post-training pruning section with the fine-tuning section, which already targets 70% sparsity, instead of the default 50%. --- docs/src/tutorials/mnist_pruning.ipynb | 52 ++++---------------------- 1 file changed, 7 insertions(+), 45 deletions(-) diff --git a/docs/src/tutorials/mnist_pruning.ipynb b/docs/src/tutorials/mnist_pruning.ipynb index cac28aa..6b1dad6 100644 --- a/docs/src/tutorials/mnist_pruning.ipynb +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -148,7 +148,7 @@ "source": [ "### Model definition", "\n", - "A simple CNN with a single Conv2d \u2192 ReLU \u2192 MaxPool block, followed by Flatten and a Linear classifier." + "A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier." ], "id": "cell-10" }, @@ -351,7 +351,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "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 \u2014 no calibration data or fine-tuning is required.\n", + "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.\n", "\n", "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." ], @@ -362,19 +362,13 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "from coreai_opt.pruning import MagnitudePruner, MagnitudePrunerConfig\n", - "\n", - "example_inputs = (torch.randn(1, 1, 28, 28),)" - ], + "source": "from coreai_opt.pruning import (\n MagnitudePruner,\n MagnitudePrunerConfig,\n ModuleMagnitudePrunerConfig,\n PruningSpec,\n)\n\nexample_inputs = (torch.randn(1, 1, 28, 28),)", "id": "cell-24" }, { "cell_type": "markdown", "metadata": {}, - "source": [ - "For this tutorial, we'll apply the default `MagnitudePrunerConfig()`, which prunes 50% of every supported weight tensor, unstructured (individual elements, not whole channels). Refer to the [Pruning Config](../pruning/config.html) page for all options." - ], + "source": "For this tutorial, we'll apply 70% unstructured magnitude pruning (individual elements, not whole channels) via `PruningSpec`. Refer to the [Pruning Config](../pruning/config.html) page for all options.", "id": "cell-25" }, { @@ -382,25 +376,13 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", - "post_training_model.eval()\n", - "\n", - "# Default config: 50% unstructured magnitude pruning on every supported weight.\n", - "post_training_config = MagnitudePrunerConfig()\n", - "\n", - "post_training_pruner = MagnitudePruner(post_training_model, post_training_config)\n", - "post_training_prepared = post_training_pruner.prepare(example_inputs)\n", - "print(\"Prepared post-training magnitude pruning at 50% sparsity\")" - ], + "source": "post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\npost_training_model.eval()\n\n# 70% unstructured magnitude pruning on every supported weight.\npost_training_config = MagnitudePrunerConfig(\n global_config=ModuleMagnitudePrunerConfig(\n op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n ),\n)\n\npost_training_pruner = MagnitudePruner(post_training_model, post_training_config)\npost_training_prepared = post_training_pruner.prepare(example_inputs)\nprint(\"Prepared post-training magnitude pruning at 70% sparsity\")", "id": "cell-26" }, { "cell_type": "markdown", "metadata": {}, - "source": [ - "After calling `prepare()`, 50% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately." - ], + "source": "After calling `prepare()`, 70% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately.", "id": "cell-27" }, { @@ -437,27 +419,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "from coreai_opt.pruning import ModuleMagnitudePrunerConfig, PruningSpec\n", - "from coreai_opt.pruning.config import PolynomialDecaySchedule\n", - "\n", - "FINE_TUNE_EPOCHS = 5\n", - "\n", - "scheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", - "\n", - "scheduled_config = MagnitudePrunerConfig(\n", - " global_config=ModuleMagnitudePrunerConfig(\n", - " op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n", - " sparsity_schedule=PolynomialDecaySchedule(\n", - " begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0\n", - " ),\n", - " ),\n", - ")\n", - "\n", - "scheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)\n", - "scheduled_prepared = scheduled_pruner.prepare(example_inputs)\n", - "print(f\"Prepared scheduled magnitude pruning targeting 70% sparsity over {FINE_TUNE_EPOCHS} epochs\")" - ], + "source": "from coreai_opt.pruning.config import PolynomialDecaySchedule\n\nFINE_TUNE_EPOCHS = 5\n\nscheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n\nscheduled_config = MagnitudePrunerConfig(\n global_config=ModuleMagnitudePrunerConfig(\n op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n sparsity_schedule=PolynomialDecaySchedule(\n begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0\n ),\n ),\n)\n\nscheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)\nscheduled_prepared = scheduled_pruner.prepare(example_inputs)\nprint(f\"Prepared scheduled magnitude pruning targeting 70% sparsity over {FINE_TUNE_EPOCHS} epochs\")", "id": "cell-31" }, { From 85c8f95bf3f91abf1df5280e6dd20be3426811c3 Mon Sep 17 00:00:00 2001 From: Utkarsh Simha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:27:00 -0700 Subject: [PATCH 3/5] docs: use 50% sparsity in MNIST pruning tutorial, bake in outputs Aligns post-training and fine-tuning sparsity at 50% and ships the notebook with executed cell outputs, showing the actual effect of pruning on accuracy end-to-end. --- docs/src/tutorials/mnist_pruning.ipynb | 434 ++++++++++++++++++------- 1 file changed, 311 insertions(+), 123 deletions(-) diff --git a/docs/src/tutorials/mnist_pruning.ipynb b/docs/src/tutorials/mnist_pruning.ipynb index 6b1dad6..6306c61 100644 --- a/docs/src/tutorials/mnist_pruning.ipynb +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -2,14 +2,15 @@ "cells": [ { "cell_type": "markdown", + "id": "cell-00", "metadata": {}, "source": [ "# Applying magnitude pruning to an MNIST model" - ], - "id": "cell-00" + ] }, { "cell_type": "markdown", + "id": "cell-01", "metadata": {}, "source": [ "In this tutorial, we will be providing a basic introduction to pruning a model with CoreAI-Opt.\n", @@ -29,30 +30,30 @@ "- [Post-Training Magnitude Pruning](#Post-Training-Magnitude-Pruning)\n", "- [Magnitude Pruning with Fine-Tuning](#Magnitude-Pruning-with-Fine-Tuning)\n", "- [Export to Core AI](#Export-to-Core-AI)" - ], - "id": "cell-01" + ] }, { "cell_type": "markdown", + "id": "cell-02", "metadata": {}, "source": [ "## Setup" - ], - "id": "cell-02" + ] }, { "cell_type": "markdown", + "id": "cell-03", "metadata": {}, "source": [ "We will be using a basic CNN model and train it on the MNIST dataset and observe its final accuracy.\n", "\n", "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." - ], - "id": "cell-03" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, + "id": "cell-04", "metadata": {}, "outputs": [], "source": [ @@ -65,12 +66,12 @@ "from torch.utils.data import DataLoader\n", "from torchinfo import summary\n", "from torchvision import datasets, transforms" - ], - "id": "cell-04" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, + "id": "cell-05", "metadata": {}, "outputs": [], "source": [ @@ -79,26 +80,37 @@ " DEVICE = torch.device(\"mps\")\n", "else:\n", " DEVICE = torch.device(\"cpu\")" - ], - "id": "cell-05" + ] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 3, + "id": "cell-06", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "SEED = 1976\n", "\n", "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)" - ], - "id": "cell-06" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, + "id": "cell-07", "metadata": { "tags": [ "parameters" @@ -108,22 +120,22 @@ "source": [ "# Used to save intermediate results and datasets\n", "SAVE_DIRECTORY = \".\"" - ], - "id": "cell-07" + ] }, { "cell_type": "markdown", + "id": "cell-08", "metadata": {}, "source": [ - "### MNIST Dataset download", + "### MNIST Dataset download\n", "\n", "Helper to download the MNIST dataset with standard normalization applied." - ], - "id": "cell-08" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, + "id": "cell-09", "metadata": {}, "outputs": [], "source": [ @@ -139,22 +151,22 @@ " train = datasets.MNIST(download_path, train=True, download=True, transform=transform)\n", " test = datasets.MNIST(download_path, train=False, download=True, transform=transform)\n", " return train, test" - ], - "id": "cell-09" + ] }, { "cell_type": "markdown", + "id": "cell-10", "metadata": {}, "source": [ - "### Model definition", + "### Model definition\n", "\n", "A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier." - ], - "id": "cell-10" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, + "id": "cell-11", "metadata": {}, "outputs": [], "source": [ @@ -173,22 +185,22 @@ "\n", " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", " return self.model(x)" - ], - "id": "cell-11" + ] }, { "cell_type": "markdown", + "id": "cell-12", "metadata": {}, "source": [ - "### Training and Evaluation", + "### Training and Evaluation\n", "\n", "Standard PyTorch training loop and evaluation function that computes accuracy." - ], - "id": "cell-12" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, + "id": "cell-13", "metadata": {}, "outputs": [], "source": [ @@ -231,12 +243,12 @@ " total += ground_truth.size(0)\n", " num_correct += (predicted == ground_truth).sum().item()\n", " return num_correct / total" - ], - "id": "cell-13" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, + "id": "cell-14", "metadata": {}, "outputs": [], "source": [ @@ -246,12 +258,12 @@ "train_dataset, test_dataset = download_mnist_dataset(\n", " download_path=DOWNLOAD_PATH\n", ")" - ], - "id": "cell-14" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, + "id": "cell-15", "metadata": {}, "outputs": [], "source": [ @@ -260,51 +272,99 @@ "# Instantiate dataloaders\n", "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n", "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE)" - ], - "id": "cell-15" + ] }, { "cell_type": "markdown", + "id": "cell-16", "metadata": {}, "source": [ "The CNN model used for this tutorial contains a single Conv2d, ReLU, MaxPool, Flatten, and Linear layer. Here's the structure:" - ], - "id": "cell-16" + ] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 11, + "id": "cell-17", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "==========================================================================================\n", + "Layer (type:depth-idx) Output Shape Param #\n", + "==========================================================================================\n", + "MnistNetwork [1, 10] --\n", + "├─Sequential: 1-1 [1, 10] --\n", + "│ └─Conv2d: 2-1 [1, 12, 28, 28] 120\n", + "│ └─ReLU: 2-2 [1, 12, 28, 28] --\n", + "│ └─MaxPool2d: 2-3 [1, 12, 14, 14] --\n", + "│ └─Flatten: 2-4 [1, 2352] --\n", + "│ └─Linear: 2-5 [1, 10] 23,530\n", + "==========================================================================================\n", + "Total params: 23,650\n", + "Trainable params: 23,650\n", + "Non-trainable params: 0\n", + "Total mult-adds (Units.MEGABYTES): 0.12\n", + "==========================================================================================\n", + "Input size (MB): 0.00\n", + "Forward/backward pass size (MB): 0.08\n", + "Params size (MB): 0.09\n", + "Estimated Total Size (MB): 0.17\n", + "==========================================================================================" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ "basic_cnn_model = MnistNetwork(num_classes=10)\n", "\n", "# Print summary of model\n", "summary(basic_cnn_model, input_size=(1, 1, 28, 28))" - ], - "id": "cell-17" + ] }, { "cell_type": "markdown", + "id": "cell-18", "metadata": {}, "source": [ "### Train baseline model" - ], - "id": "cell-18" + ] }, { "cell_type": "markdown", + "id": "cell-19", "metadata": {}, "source": [ "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." - ], - "id": "cell-19" + ] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 12, + "id": "cell-20", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " Epoch 1: loss=0.3126\n", + " Epoch 2: loss=0.1170\n", + " Epoch 3: loss=0.0821\n", + " Epoch 4: loss=0.0669\n", + " Epoch 5: loss=0.0590\n", + " Epoch 6: loss=0.0522\n", + " Epoch 7: loss=0.0481\n", + " Epoch 8: loss=0.0452\n", + " Epoch 9: loss=0.0414\n", + " Epoch 10: loss=0.0379\n" + ] + } + ], "source": [ "EPOCHS = 10\n", "\n", @@ -325,116 +385,216 @@ "\n", "# Save trained weights for reuse in pruning sections\n", "pretrained_state_dict = basic_cnn_model.state_dict()" - ], - "id": "cell-20" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, + "id": "cell-21", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Baseline accuracy: 0.9798\n" + ] + } + ], "source": [ "accuracy = eval_model(basic_cnn_model, test_loader)\n", "print(f\"Baseline accuracy: {accuracy:.4f}\")" - ], - "id": "cell-21" + ] }, { "cell_type": "markdown", + "id": "cell-22", "metadata": {}, "source": [ "## Post-Training Magnitude Pruning" - ], - "id": "cell-22" + ] }, { "cell_type": "markdown", + "id": "cell-23", "metadata": {}, "source": [ "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.\n", "\n", "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." - ], - "id": "cell-23" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 14, + "id": "cell-24", "metadata": {}, - "outputs": [], - "source": "from coreai_opt.pruning import (\n MagnitudePruner,\n MagnitudePrunerConfig,\n ModuleMagnitudePrunerConfig,\n PruningSpec,\n)\n\nexample_inputs = (torch.randn(1, 1, 28, 28),)", - "id": "cell-24" + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "W0714 17:27:02.612000 83677 torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.\n", + "scikit-learn version 1.9.0 is not supported. Minimum required version: 0.17. Maximum required version: 1.5.1. Disabling scikit-learn conversion API.\n", + "Torch version 2.11.0 has not been tested with coremltools. You may run into unexpected errors. Torch 2.7.0 is the most recent version that has been tested.\n" + ] + } + ], + "source": [ + "from coreai_opt.pruning import (\n", + " MagnitudePruner,\n", + " MagnitudePrunerConfig,\n", + " ModuleMagnitudePrunerConfig,\n", + " PruningSpec,\n", + ")\n", + "\n", + "example_inputs = (torch.randn(1, 1, 28, 28),)" + ] }, { "cell_type": "markdown", + "id": "cell-25", "metadata": {}, - "source": "For this tutorial, we'll apply 70% unstructured magnitude pruning (individual elements, not whole channels) via `PruningSpec`. Refer to the [Pruning Config](../pruning/config.html) page for all options.", - "id": "cell-25" + "source": [ + "For this tutorial, we'll apply 50% unstructured magnitude pruning (individual elements, not whole channels) via `PruningSpec`. Refer to the [Pruning Config](../pruning/config.html) page for all options." + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 15, + "id": "cell-26", "metadata": {}, - "outputs": [], - "source": "post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\npost_training_model.eval()\n\n# 70% unstructured magnitude pruning on every supported weight.\npost_training_config = MagnitudePrunerConfig(\n global_config=ModuleMagnitudePrunerConfig(\n op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n ),\n)\n\npost_training_pruner = MagnitudePruner(post_training_model, post_training_config)\npost_training_prepared = post_training_pruner.prepare(example_inputs)\nprint(\"Prepared post-training magnitude pruning at 70% sparsity\")", - "id": "cell-26" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Prepared post-training magnitude pruning at 50% sparsity\n" + ] + } + ], + "source": [ + "post_training_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", + "post_training_model.eval()\n", + "\n", + "# 50% unstructured magnitude pruning on every supported weight.\n", + "post_training_config = MagnitudePrunerConfig(\n", + " global_config=ModuleMagnitudePrunerConfig(\n", + " op_state_spec={\"weight\": PruningSpec(target_sparsity=0.5)},\n", + " ),\n", + ")\n", + "\n", + "post_training_pruner = MagnitudePruner(post_training_model, post_training_config)\n", + "post_training_prepared = post_training_pruner.prepare(example_inputs)\n", + "print(\"Prepared post-training magnitude pruning at 50% sparsity\")" + ] }, { "cell_type": "markdown", + "id": "cell-27", "metadata": {}, - "source": "After calling `prepare()`, 70% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately.", - "id": "cell-27" + "source": [ + "After calling `prepare()`, 50% of the values in each supported weight tensor are already zeroed out, so we can measure the accuracy impact immediately." + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 16, + "id": "cell-28", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Post-training pruning accuracy: 0.8805\n" + ] + } + ], "source": [ "post_training_accuracy = eval_model(post_training_prepared, test_loader)\n", "print(f\"Post-training pruning accuracy: {post_training_accuracy:.4f}\")" - ], - "id": "cell-28" + ] }, { "cell_type": "markdown", + "id": "cell-29", "metadata": {}, "source": [ "## Magnitude Pruning with Fine-Tuning" - ], - "id": "cell-29" + ] }, { "cell_type": "markdown", + "id": "cell-30", "metadata": {}, "source": [ "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.\n", "\n", - "Here we target 70% sparsity, ramped in via a `PolynomialDecaySchedule` over the 5 epochs of fine-tuning." - ], - "id": "cell-30" + "Here we target 50% sparsity, ramped in via a `PolynomialDecaySchedule` over the 5 epochs of fine-tuning." + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 17, + "id": "cell-31", "metadata": {}, - "outputs": [], - "source": "from coreai_opt.pruning.config import PolynomialDecaySchedule\n\nFINE_TUNE_EPOCHS = 5\n\nscheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n\nscheduled_config = MagnitudePrunerConfig(\n global_config=ModuleMagnitudePrunerConfig(\n op_state_spec={\"weight\": PruningSpec(target_sparsity=0.7)},\n sparsity_schedule=PolynomialDecaySchedule(\n begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0\n ),\n ),\n)\n\nscheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)\nscheduled_prepared = scheduled_pruner.prepare(example_inputs)\nprint(f\"Prepared scheduled magnitude pruning targeting 70% sparsity over {FINE_TUNE_EPOCHS} epochs\")", - "id": "cell-31" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Prepared scheduled magnitude pruning targeting 50% sparsity over 5 epochs\n" + ] + } + ], + "source": [ + "from coreai_opt.pruning.config import PolynomialDecaySchedule\n", + "\n", + "FINE_TUNE_EPOCHS = 5\n", + "\n", + "scheduled_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", + "\n", + "scheduled_config = MagnitudePrunerConfig(\n", + " global_config=ModuleMagnitudePrunerConfig(\n", + " op_state_spec={\"weight\": PruningSpec(target_sparsity=0.5)},\n", + " sparsity_schedule=PolynomialDecaySchedule(\n", + " begin_step=0, total_iters=FINE_TUNE_EPOCHS, power=3.0\n", + " ),\n", + " ),\n", + ")\n", + "\n", + "scheduled_pruner = MagnitudePruner(scheduled_model, scheduled_config)\n", + "scheduled_prepared = scheduled_pruner.prepare(example_inputs)\n", + "print(f\"Prepared scheduled magnitude pruning targeting 50% sparsity over {FINE_TUNE_EPOCHS} epochs\")" + ] }, { "cell_type": "markdown", + "id": "cell-32", "metadata": {}, "source": [ "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." - ], - "id": "cell-32" + ] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 18, + "id": "cell-33", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " Epoch 1: loss=0.0299, accuracy=0.9750\n", + " Epoch 2: loss=0.0354, accuracy=0.9530\n", + " Epoch 3: loss=0.0619, accuracy=0.9727\n", + " Epoch 4: loss=0.0561, accuracy=0.9793\n", + " Epoch 5: loss=0.0493, accuracy=0.9796\n" + ] + } + ], "source": [ "fine_tune_optimizer = torch.optim.SGD(scheduled_prepared.parameters(), lr=1e-3)\n", "\n", @@ -454,52 +614,60 @@ "\n", "scheduled_prepared = scheduled_prepared.eval()\n", "print(\"\\n\".join(fine_tune_results))" - ], - "id": "cell-33" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 19, + "id": "cell-34", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Scheduled magnitude pruning accuracy (50% sparsity, fine-tuned): 0.9796\n" + ] + } + ], "source": [ "scheduled_accuracy = eval_model(scheduled_prepared, test_loader)\n", - "print(f\"Scheduled magnitude pruning accuracy (70% sparsity, fine-tuned): {scheduled_accuracy:.4f}\")" - ], - "id": "cell-34" + "print(f\"Scheduled magnitude pruning accuracy (50% sparsity, fine-tuned): {scheduled_accuracy:.4f}\")" + ] }, { "cell_type": "markdown", + "id": "cell-35", "metadata": {}, "source": [ "## Export to Core AI" - ], - "id": "cell-35" + ] }, { "cell_type": "markdown", + "id": "cell-36", "metadata": {}, "source": [ "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`.\n", "\n", "We'll export the fine-tuned, scheduled-pruning model from the previous section." - ], - "id": "cell-36" + ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 20, + "id": "cell-37", "metadata": {}, "outputs": [], "source": [ "from coreai_opt import ExportBackend\n", "\n", "coreai_model = scheduled_pruner.finalize(backend=ExportBackend.CoreAI)" - ], - "id": "cell-37" + ] }, { "cell_type": "markdown", + "id": "cell-38", "metadata": {}, "source": [ "The export proceeds in three steps:\n", @@ -507,14 +675,35 @@ "- Trace the model with `torch.export.export()` to obtain a graph representation.\n", "- Apply `cast_to_16_bit_precision()` to cast remaining FP32 parameters to FP16 for optimal on-device performance.\n", "- Convert the exported program to Core AI format using `coreai-torch.TorchConverter`." - ], - "id": "cell-38" + ] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], + "execution_count": 21, + "id": "cell-39", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
coreai-torch 0.4.1: converting 1 program(s) to Core AI\n",
+       "
\n" + ], + "text/plain": [ + "\u001b[1;36mcoreai-torch\u001b[0m \u001b[1;2;36m0.4\u001b[0m\u001b[2m.\u001b[0m\u001b[1;2;36m1\u001b[0m: converting \u001b[1;36m1\u001b[0m \u001b[1;35mprogram\u001b[0m\u001b[1m(\u001b[0ms\u001b[1m)\u001b[0m to Core AI\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Exported: exported_model.aimodel\n" + ] + } + ], "source": [ "import shutil\n", "\n", @@ -533,13 +722,12 @@ " shutil.rmtree(output_path)\n", "coreai_program.save_asset(output_path)\n", "print(f\"Exported: {output_path}\")" - ], - "id": "cell-39" + ] } ], "metadata": { "kernelspec": { - "display_name": ".venv-tutorial", + "display_name": ".venv (3.11.13)", "language": "python", "name": "python3" }, @@ -553,7 +741,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.14" + "version": "3.11.13" } }, "nbformat": 4, From 90a504e9c7cb9877ea9267b9223e71e40c3bf84e Mon Sep 17 00:00:00 2001 From: Utkarsh Simha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:27:01 -0700 Subject: [PATCH 4/5] docs: strip coremltools import warning from pruning tutorial output Removes a local scikit-learn/coremltools version warning baked into the notebook's output; it's environment noise, not tutorial content. --- docs/src/tutorials/mnist_pruning.ipynb | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/docs/src/tutorials/mnist_pruning.ipynb b/docs/src/tutorials/mnist_pruning.ipynb index 6306c61..ab29976 100644 --- a/docs/src/tutorials/mnist_pruning.ipynb +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -429,17 +429,7 @@ "execution_count": 14, "id": "cell-24", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "W0714 17:27:02.612000 83677 torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.\n", - "scikit-learn version 1.9.0 is not supported. Minimum required version: 0.17. Maximum required version: 1.5.1. Disabling scikit-learn conversion API.\n", - "Torch version 2.11.0 has not been tested with coremltools. You may run into unexpected errors. Torch 2.7.0 is the most recent version that has been tested.\n" - ] - } - ], + "outputs": [], "source": [ "from coreai_opt.pruning import (\n", " MagnitudePruner,\n", From dbb2e2b35468ae26149ec92336ba4fbbacc87f5d Mon Sep 17 00:00:00 2001 From: Utkarsh Simha <135899523+u-simha@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:27:02 -0700 Subject: [PATCH 5/5] docs: fix broken pruning config link in MNIST pruning tutorial --- docs/src/tutorials/mnist_pruning.ipynb | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/src/tutorials/mnist_pruning.ipynb b/docs/src/tutorials/mnist_pruning.ipynb index ab29976..6788ca2 100644 --- a/docs/src/tutorials/mnist_pruning.ipynb +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -445,9 +445,7 @@ "cell_type": "markdown", "id": "cell-25", "metadata": {}, - "source": [ - "For this tutorial, we'll apply 50% unstructured magnitude pruning (individual elements, not whole channels) via `PruningSpec`. Refer to the [Pruning Config](../pruning/config.html) page for all options." - ] + "source": "For this tutorial, we'll apply 50% unstructured magnitude pruning (individual elements, not whole channels) via `PruningSpec`. Refer to the [Pruning Config](https://apple.github.io/coreai-optimization/pruning/config.html) page for all options." }, { "cell_type": "code",