In this tutorial, we build a complete end-to-end pipeline using NVIDIA Model Optimizer to train, prune, and fine-tune a deep learning model directly in Google Colab. We start by setting up the environment and preparing the CIFAR-10 dataset, then define a ResNet architecture and train it to establish a strong baseline. From there, we apply FastNAS pruning to systematically reduce the model’s complexity under FLOPs constraints while preserving performance. We also handle real-world compatibility issues, restore the optimized subnet, and fine-tune it to recover accuracy. By the end, we have a fully working workflow that takes a model from training to deployment-ready optimization, all within a single streamlined setup. Check out the Full Implementation Coding Notebook. Copy CodeCopiedUse a different Browser !pip -q install -U nvidia-modelopt torchvision torchprofile tqdm import math import os import random import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader, Subset from torchvision.models.resnet import BasicBlock from tqdm.auto import tqdm import modelopt.torch.opt as mto import modelopt.torch.prune as mtp SEED = 123 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED) FAST_MODE = True batch_size = 256 if FAST_MODE else 512 baseline_epochs = 20 if FAST_MODE else 120 finetune_epochs = 12 if FAST_MODE else 120 train_subset_size = 12000 if FAST_MODE else None val_subset_size = 2000 if FAST_MODE else None test_subset_size = 4000 if FAST_MODE else None target_flops = 60e6 We begin by installing all required dependencies and importing the necessary libraries to set up our environment. We initialize seeds to ensure reproducibility and configure the device to leverage a GPU if available. We also define key runtime parameters, such as batch size, epochs, dataset subsets, and FLOP constraints, to control the overall experiment. Copy CodeCopiedUse a different Browser def seed_worker(worker_id): worker_seed = SEED + worker_id np.random.seed(worker_seed) random.seed(worker_seed) def build_cifar10_loaders(train_batch_size=256, train_subset_size=None, val_subset_size=None, test_subset_size=None): normalize = transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2470, 0.2435, 0.2616], ) train_transform = transforms.Compose([ transforms.ToTensor(), transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, padding=4), normalize, ]) eval_transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) train_full = torchvision.datasets.CIFAR10( root=”./data”, train=True, transform=train_transform, download=True ) val_full = torchvision.datasets.CIFAR10( root=”./data”, train=True, transform=eval_transform, download=True ) test_full = torchvision.datasets.CIFAR10( root=”./data”, train=False, transform=eval_transform, download=True ) n_trainval = len(train_full) ids = np.arange(n_trainval) np.random.shuffle(ids) n_train = int(n_trainval * 0.9) train_ids = ids[:n_train] val_ids = ids[n_train:] if train_subset_size is not None: train_ids = train_ids[:min(train_subset_size, len(train_ids))] if val_subset_size is not None: val_ids = val_ids[:min(val_subset_size, len(val_ids))] test_ids = np.arange(len(test_full)) if test_subset_size is not None: test_ids = test_ids[:min(test_subset_size, len(test_ids))] train_ds = Subset(train_full, train_ids.tolist()) val_ds = Subset(val_full, val_ids.tolist()) test_ds = Subset(test_full, test_ids.tolist()) num_workers = min(2, os.cpu_count() or 1) g = torch.Generator() g.manual_seed(SEED) train_loader = DataLoader( train_ds, batch_size=train_batch_size, shuffle=True, num_workers=num_workers, pin_memory=torch.cuda.is_available(), worker_init_fn=seed_worker, generator=g, ) val_loader = DataLoader( val_ds, batch_size=512, shuffle=False, num_workers=num_workers, pin_memory=torch.cuda.is_available(), worker_init_fn=seed_worker, ) test_loader = DataLoader( test_ds, batch_size=512, shuffle=False, num_workers=num_workers, pin_memory=torch.cuda.is_available(), worker_init_fn=seed_worker, ) print(f”Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}”) return train_loader, val_loader, test_loader train_loader, val_loader, test_loader = build_cifar10_loaders( train_batch_size=batch_size, train_subset_size=train_subset_size, val_subset_size=val_subset_size, test_subset_size=test_subset_size, ) We construct the full data pipeline by preparing CIFAR-10 datasets with appropriate augmentations and normalization. We split the dataset to reduce its size and speed up experimentation. We then create efficient data loaders that ensure proper batching, shuffling, and reproducible data handling. Copy CodeCopiedUse a different Browser def _weights_init(m): if isinstance(m, (nn.Linear, nn.Conv2d)): nn.init.kaiming_normal_(m.weight) class LambdaLayer(nn.Module): def __init__(self, lambd): super().__init__() self.lambd = lambd def forward(self, x): return self.lambd(x) class ResNet(nn.Module): def __init__(self, num_blocks, num_classes=10): super().__init__() self.in_planes = 16 self.layers = nn.Sequential( nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False), nn.BatchNorm2d(16), nn.ReLU(), self._make_layer(16, num_blocks, stride=1), self._make_layer(32, num_blocks, stride=2), self._make_layer(64, num_blocks, stride=2), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(64, num_classes), ) self.apply(_weights_init) def _make_layer(self, planes, num_blocks, stride): strides = [stride] + [1] * (num_blocks – 1) layers = [] for s in strides: downsample = None if s != 1 or self.in_planes != planes: downsample = LambdaLayer( lambda x: F.pad( x[:, :, ::2, ::2], (0, 0, 0, 0, planes // 4, planes // 4), “constant”, 0, ) ) layers.append(BasicBlock(self.in_planes, planes, s, downsample)) self.in_planes = planes return nn.Sequential(*layers) def forward(self, x): return self.layers(x) def resnet20(): return ResNet(num_blocks=3).to(device) We define the ResNet20 architecture from scratch, including custom initialization and shortcut handling through lambda layers. We structure the network using convolutional blocks and residual connections to capture hierarchical features. We finally encapsulate the model creation into a reusable function that moves it directly to the selected device. Copy CodeCopiedUse a different Browser class CosineLRwithWarmup(torch.optim.lr_scheduler._LRScheduler): def __init__(self, optimizer, warmup_steps, decay_steps, warmup_lr=0.0, last_epoch=-1): self.warmup_steps = warmup_steps self.warmup_lr = warmup_lr self.decay_steps = max(decay_steps, 1) super().__init__(optimizer, last_epoch) def get_lr(self): if self.last_epoch < self.warmup_steps: return [ (base_lr – self.warmup_lr) * self.last_epoch / max(self.warmup_steps, 1) + self.warmup_lr for base_lr in self.base_lrs ] current_steps = self.last_epoch – self.warmup_steps return [ 0.5 * base_lr * (1 + math.cos(math.pi * current_steps / self.decay_steps)) for base_lr in self.base_lrs ] def get_optimizer_scheduler(model, lr, weight_decay, warmup_steps, decay_steps): optimizer = torch.optim.SGD( filter(lambda p: p.requires_grad, model.parameters()), lr=lr, momentum=0.9, weight_decay=weight_decay, ) scheduler = CosineLRwithWarmup(optimizer, warmup_steps, decay_steps) return optimizer, scheduler def loss_fn_default(model, outputs, labels): return F.cross_entropy(outputs, labels) def train_one_epoch(model, loader, optimizer, scheduler, loss_fn=loss_fn_default): model.train() running_loss = 0.0 total = 0 for images, labels in loader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) outputs = model(images) loss = loss_fn(model, outputs, labels) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() scheduler.step() running_loss += loss.item() * labels.size(0) total += labels.size(0) return running_loss / max(total, 1) @torch.no_grad() def evaluate(model, loader): model.eval() correct = 0 total = 0 for images, labels in loader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) logits = model(images) preds = logits.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) return 100.0 * correct / max(total, 1) def train_model(model, train_loader, val_loader, epochs, ckpt_path, lr=None, weight_decay=1e-4, print_every=1): if lr is None: lr = 0.1 * batch_size / 128 steps_per_epoch = len(train_loader) warmup_steps = max(1, 2 * steps_per_epoch if FAST_MODE else 5 * steps_per_epoch) decay_steps = max(1, epochs * steps_per_epoch) optimizer, scheduler = get_optimizer_scheduler( model=model, lr=lr, weight_decay=weight_decay, warmup_steps=warmup_steps, decay_steps=decay_steps, ) best_val = -1.0 best_epoch = -1 print(f”Training for {epochs} epochs…”) for epoch in tqdm(range(1, epochs + 1)): train_loss = train_one_epoch(model, train_loader, optimizer, scheduler) val_acc = evaluate(model, val_loader) if val_acc >= best_val: best_val = val_acc best_epoch = epoch torch.save(model.state_dict(), ckpt_path) if epoch == 1