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..6788ca2 --- /dev/null +++ b/docs/src/tutorials/mnist_pruning.ipynb @@ -0,0 +1,737 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-00", + "metadata": {}, + "source": [ + "# Applying magnitude pruning to an MNIST model" + ] + }, + { + "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", + "\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)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-02", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "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." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "cell-04", + "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" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "cell-05", + "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\")" + ] + }, + { + "cell_type": "code", + "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)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "cell-07", + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [], + "source": [ + "# Used to save intermediate results and datasets\n", + "SAVE_DIRECTORY = \".\"" + ] + }, + { + "cell_type": "markdown", + "id": "cell-08", + "metadata": {}, + "source": [ + "### MNIST Dataset download\n", + "\n", + "Helper to download the MNIST dataset with standard normalization applied." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "cell-09", + "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" + ] + }, + { + "cell_type": "markdown", + "id": "cell-10", + "metadata": {}, + "source": [ + "### Model definition\n", + "\n", + "A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "cell-11", + "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)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-12", + "metadata": {}, + "source": [ + "### Training and Evaluation\n", + "\n", + "Standard PyTorch training loop and evaluation function that computes accuracy." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "cell-13", + "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" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "cell-14", + "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", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "cell-15", + "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)" + ] + }, + { + "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:" + ] + }, + { + "cell_type": "code", + "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))" + ] + }, + { + "cell_type": "markdown", + "id": "cell-18", + "metadata": {}, + "source": [ + "### Train baseline model" + ] + }, + { + "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." + ] + }, + { + "cell_type": "code", + "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", + "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()" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "cell-21", + "metadata": {}, + "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}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-22", + "metadata": {}, + "source": [ + "## Post-Training Magnitude Pruning" + ] + }, + { + "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." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "cell-24", + "metadata": {}, + "outputs": [], + "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 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", + "execution_count": 15, + "id": "cell-26", + "metadata": {}, + "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()`, 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": 16, + "id": "cell-28", + "metadata": {}, + "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}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-29", + "metadata": {}, + "source": [ + "## Magnitude Pruning with Fine-Tuning" + ] + }, + { + "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 50% sparsity, ramped in via a `PolynomialDecaySchedule` over the 5 epochs of fine-tuning." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "cell-31", + "metadata": {}, + "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." + ] + }, + { + "cell_type": "code", + "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", + "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))" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "cell-34", + "metadata": {}, + "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 (50% sparsity, fine-tuned): {scheduled_accuracy:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-35", + "metadata": {}, + "source": [ + "## Export to Core AI" + ] + }, + { + "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." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "cell-37", + "metadata": {}, + "outputs": [], + "source": [ + "from coreai_opt import ExportBackend\n", + "\n", + "coreai_model = scheduled_pruner.finalize(backend=ExportBackend.CoreAI)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-38", + "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`." + ] + }, + { + "cell_type": "code", + "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", + "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}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv (3.11.13)", + "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.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}