φ-Register: φ-Pruning on GPU

Runtime → Change runtime type → T4 GPU (if not already set)

This notebook runs the phi-prune quickstart pipeline with GPU acceleration: dense fine-tune → prune → verify → mask-aware fine-tune → Fibonacci encoding → export. CIFAR-10 stays on the GPU and is resized there, and training and evaluation run under fp16 autocast in channels-last layout: the T4's tensor cores (65 TFLOPS) take fp16, while fp32 runs on its CUDA cores (8 TFLOPS).

0. Install

# Option A: pip install straight from GitHub
# !pip install -q git+https://github.com/ezexe/phi.git

# Option B: Upload the package zip, then:
# from google.colab import files
# uploaded = files.upload()  # upload phi-prune.zip
# !unzip -q phi-prune.zip -d phi-prune
# !pip install -q phi-prune/

# Option C: Clone from GitHub; a re-run pulls new commits into the existing clone.
# A regular install is importable in this running kernel. An editable (-e) install registers the package
# through a .pth file, which Python reads only at startup, so the import in Setup would fail until the
# runtime restarts.
!test -d phi || git clone -q --depth 1 https://github.com/ezexe/phi.git
!git -C phi pull -q --ff-only
!pip install -q ./phi

import torch
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    props = torch.cuda.get_device_properties(0)
    print(f"GPU: {props.name}")
    print(f"VRAM: {props.total_memory / 1e9:.1f} GB")
else:
    print("No GPU: Runtime → Change runtime type → T4 GPU, then run all cells again")

1. Setup

Imports and a check that the installed package has finetune(amp=...), config knobs, then the GPU evaluate helper.

import inspect

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from phi_prune import prune, finetune, check
from phi_prune.encoding import FibonacciEncoder
from phi_prune.export import save_checkpoint

if "amp" not in inspect.signature(finetune).parameters:
    raise RuntimeError("phi-prune 0.2.0+ is needed for finetune(amp=...): rerun the install cell, "
                       "then restart the runtime")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
USE_AMP = device.type == "cuda"        # fp16 autocast + loss scaling; T4 tensor cores have no bf16/TF32
torch.backends.cudnn.benchmark = True  # fixed input shapes, so cuDNN's per-shape autotuning pays off
torch.manual_seed(0)

IMG_SIZE = 224                    # ImageNet weights expect ~224 px; 160/128 is faster, less accurate
TRAIN_BATCH = 128                 # as in the original notebook; rescale the learning rates if you change it
EVAL_BATCH = 512                  # eval keeps no activations for backward, so its batches can be larger
DENSE_EPOCHS, DENSE_LR = 2, 0.01  # trains the new 10-class head (and the backbone) before pruning
FT_EPOCHS, FT_LR = 5, 0.001       # mask-aware recovery after pruning, as in the quickstart
print(f"Device: {device}, AMP: {USE_AMP}")

try:
    from phi_prune import evaluate as package_evaluate
except ImportError:
    raise RuntimeError("phi-prune 0.2.2+ is needed for evaluate(): rerun the install cell, "
                       "then restart the runtime") from None


def evaluate(model, loader):
    """Top-1 % under fp16 autocast on the GPU. A batch whose fp16 logits overflow to inf/NaN is
    re-run in fp32: argmax over NaN picks class 0, which would score chance level (10.00%)."""
    return package_evaluate(model, loader, device, amp=USE_AMP)

2. Load model & data

# Pretrained ResNet-18, swap head for CIFAR-10
model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(512, 10)
model = model.to(device, memory_format=torch.channels_last)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")

# CIFAR-10 stays on the GPU as uint8 (~184 MB) and each batch is normalized, flipped and resized there.
# This replaces the per-image PIL Resize/ToTensor/Normalize that ran in 2 DataLoader workers, and the
# host-to-device copy of every batch at 224 px float32 (77 MB per 128 images).
MEAN = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)


class GPULoader:
    def __init__(self, dataset, batch_size, train):
        self.x = torch.from_numpy(dataset.data).to(device).permute(0, 3, 1, 2).contiguous()  # N,3,32,32
        self.y = torch.as_tensor(dataset.targets, device=device)
        self.batch_size, self.train = batch_size, train

    def __len__(self):
        return -(-len(self.y) // self.batch_size)

    def __iter__(self):
        n = len(self.y)
        order = torch.randperm(n, device=device) if self.train else torch.arange(n, device=device)
        for i in range(0, n, self.batch_size):
            idx = order[i:i + self.batch_size]
            x = self.x[idx].float().div_(255).sub_(MEAN).div_(STD)
            if self.train:  # random horizontal flip
                flip = torch.rand(len(idx), device=device) < 0.5
                x = torch.where(flip.view(-1, 1, 1, 1), x.flip(3), x)
            x = F.interpolate(x, size=(IMG_SIZE, IMG_SIZE), mode="bilinear", align_corners=False)
            yield x.contiguous(memory_format=torch.channels_last), self.y[idx]


trainset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True)
testset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True)
train_loader = GPULoader(trainset, TRAIN_BATCH, train=True)
test_loader = GPULoader(testset, EVAL_BATCH, train=False)
print(f"Train: {len(trainset)}, Test: {len(testset)}")

3. Dense baseline

The pretrained backbone gets a new, randomly initialized 10-class fc, so the dense model scores ~10% (chance) until that head is trained. Pruning is measured against the dense model after DENSE_EPOCHS of fine-tuning.

%%time
dense_results = finetune(model, train_loader, epochs=DENSE_EPOCHS, lr=DENSE_LR, device=device,
                         val_loader=test_loader, amp=USE_AMP)
dense_acc = dense_results["best_val_acc"]
print(f"Dense: {dense_acc:.2f}%")

4. Prune & verify

Only conv layers are pruned, so the 10-class fc head stays dense. Pruning fc.weight along its 10 output rows would keep at most 5 of them (no two adjacent), leaving the other classes with a bias-only, input-independent logit, which caps test accuracy at 60%. From phi-prune 0.2.0 on, prune() leaves the head dense by default (prune_head=False); layer_types=(nn.Conv2d,) makes the conv-only choice explicit. The ResNet-20 experiment in .docs/experiment/Zeckendorf.py prunes conv layers only, too.

pruned_model, masks = prune(model, inplace=False, layer_types=(nn.Conv2d,))
stats = pruned_model._phi_prune_stats
print(f"Density (conv weights): {stats['density']:.1%}")
print(f"Pruned layers: {stats['pruned_layers']}")

report = check(pruned_model, masks)
print(f"All masks valid: {report['_summary']['all_masks_valid']}")
print(f"All zeros enforced: {report['_summary']['all_zeros_enforced']}")

# Accuracy after pruning, before fine-tune
pruned_acc = evaluate(pruned_model, test_loader)
print(f"Pruned (no fine-tune): {pruned_acc:.2f}%")

5. Fine-tune

finetune(..., amp=USE_AMP) runs the package's mask-aware fine-tune under fp16 autocast with loss scaling on the GPU. Raise FT_EPOCHS for better recovery.

%%time
ft_results = finetune(pruned_model, train_loader, epochs=FT_EPOCHS, masks=masks, lr=FT_LR, device=device,
                      val_loader=test_loader, amp=USE_AMP)
print(f"Best val accuracy: {ft_results['best_val_acc']:.2f}%")

6. Fibonacci encoding

encoder = FibonacciEncoder(n_digits=10)  # 144 levels
print(f"Grid: {encoder.n_levels} levels, max value: {encoder.max_value}")

sample_param = next(n for n, p in pruned_model.named_parameters() if n in masks)
param = dict(pruned_model.named_parameters())[sample_param]
encoded, scale, rmse = encoder.encode_tensor(param.data, mask=masks[sample_param], axis=0)  # per channel
print(f"Sample layer ({sample_param}): RMSE = {rmse:.6f}")

7. Summary & export

report = check(pruned_model, masks)  # re-verify after fine-tuning, not only right after pruning
rows = [
    (f"Dense ({DENSE_EPOCHS} ep ft)", f"{dense_acc:.2f}%"),
    ("Pruned (no ft)", f"{pruned_acc:.2f}%"),
    (f"Pruned ({FT_EPOCHS} ep ft)", f"{ft_results['best_val_acc']:.2f}%"),
    ("Drop vs dense", f"{dense_acc - ft_results['best_val_acc']:.2f} pts"),
    ("Density (conv weights)", f"{stats['density']:.1%}"),
    ("Masks valid", report["_summary"]["all_masks_valid"]),
    ("Zeros enforced", report["_summary"]["all_zeros_enforced"]),
    ("Encoding levels", encoder.n_levels),
]
print("=" * 50)
for label, value in rows:
    print(f"  {label + ':':<24}{value}")
print("=" * 50)

save_checkpoint(pruned_model, masks, "model_pruned.pt", encoder=encoder, metadata={
    "dense_acc": dense_acc, "pruned_acc": pruned_acc,
    "finetuned_acc": ft_results["best_val_acc"], "img_size": IMG_SIZE,
})
print("Saved model_pruned.pt")

# Download checkpoint locally
from google.colab import files
files.download('model_pruned.pt')

More ResNets and weights

Each cell below runs sections 3–5 (dense fine-tune, conv-only prune, mask-aware fine-tune) for one more torchvision ResNet and weight version on the CIFAR-10 already on the GPU, and prints its own framed summary. Run the helpers cell first, then any variant cells in any order, then the comparison. IMAGENET1K_V2 weights come from torchvision's newer training recipe; ResNet-18 and ResNet-34 have only IMAGENET1K_V1. ResNet-18 took about a minute per epoch on a T4, and time grows roughly with the GFLOPs each cell prints (1.8 for ResNet-18, 3.7 for ResNet-34, 4.1 for ResNet-50, 7.8 for ResNet-101, 11.5 for ResNet-152), so the larger variants take half an hour or more each. Colab disconnects a session left idle too long and ends every session within 12 hours, so run the big variants while you are around.

#@title Helpers for the variant cells below
import copy
import inspect
import json
import os
import time

import phi_prune

try:
    from google.colab import files as colab_files
except ImportError:  # outside Colab the checkpoints just stay on disk
    colab_files = None

if tuple(int(x) for x in phi_prune.__version__.split(".")[:3]) < (0, 2, 1):
    raise RuntimeError(f"phi-prune {phi_prune.__version__} predates 0.2.1's pruning fix: "
                       "rerun the install cell, then restart the runtime")

CHECKPOINT_DIR = "checkpoints"
DOWNLOAD_CHECKPOINTS = True  # Colab's disk goes with the session: download each checkpoint as it's saved
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
variant_results = {}


def build_model(arch, weights_name):
    """A torchvision ResNet with its ImageNet weights and a new 10-class fc, channels-last on the device."""
    weights = torchvision.models.get_model_weights(arch)[weights_name]
    net = torchvision.models.get_model(arch, weights=weights)
    net.fc = nn.Linear(net.fc.in_features, 10)
    return net.to(device, memory_format=torch.channels_last), weights


def run_variant(arch, weights_name, batch=TRAIN_BATCH, pattern="phi"):
    """Sections 3-5 for one variant, then its checkpoint: dense fine-tune, conv-only prune, fine-tune.

    pattern="2:4" swaps the φ mask for the 2-of-4 baseline, which needs prune(pattern=...).
    """
    prune_kw = {}
    if pattern != "phi":
        if "pattern" not in inspect.signature(prune).parameters:
            raise RuntimeError("this phi-prune has no prune(pattern=...): "
                               "rerun the install cell, then restart the runtime")
        prune_kw["pattern"] = pattern
    label = f"{arch} / {weights_name}" + ("" if pattern == "phi" else f" / {pattern}")
    train = copy.copy(train_loader)  # shares section 2's GPU tensors; only the batch size differs
    train.batch_size = batch
    scale = batch / TRAIN_BATCH      # linear lr scaling when the batch shrinks
    net, weights = build_model(arch, weights_name)
    params_m = sum(p.numel() for p in net.parameters()) / 1e6
    gflops = weights.meta.get("_ops", float("nan"))
    # current torchvision keeps ImageNet accuracy under "_metrics"; the plain key is only a fallback
    metrics = weights.meta.get("_metrics") or weights.meta.get("metrics", {})
    top1 = metrics.get("ImageNet-1K", {}).get("acc@1", float("nan"))
    print(f"{label}: {params_m:.1f}M params, {gflops:.1f} GFLOPs, "
          f"ImageNet top-1 {top1:.2f}%")
    start = time.perf_counter()
    dense = finetune(net, train, epochs=DENSE_EPOCHS, lr=DENSE_LR * scale, device=device,
                     val_loader=test_loader, amp=USE_AMP)
    pruned, variant_masks = prune(net, inplace=False, layer_types=(nn.Conv2d,), **prune_kw)
    no_ft = evaluate(pruned, test_loader)
    ft = finetune(pruned, train, epochs=FT_EPOCHS, masks=variant_masks, lr=FT_LR * scale, device=device,
                  val_loader=test_loader, amp=USE_AMP)
    summary = check(pruned, variant_masks)["_summary"]  # validates the pattern prune() applied
    row = {
        "params_m": params_m, "gflops": gflops, "batch": batch, "pattern": pattern,
        "dense": dense["best_val_acc"], "pruned_no_ft": no_ft, "pruned_ft": ft["best_val_acc"],
        "density": pruned._phi_prune_stats["density"], "layers": pruned._phi_prune_stats["pruned_layers"],
        "masks_ok": summary["all_masks_valid"] and summary["all_zeros_enforced"],
        "minutes": (time.perf_counter() - start) / 60,
    }
    suffix = "" if pattern == "phi" else "_" + pattern.replace(":", "of")  # no ':' in file names
    path = os.path.join(CHECKPOINT_DIR, f"{arch}_{weights_name}{suffix}_pruned.pt")
    save_checkpoint(pruned, variant_masks, path, metadata={"arch": arch, "weights": weights_name, **row})
    row["checkpoint"], row["checkpoint_mb"] = path, os.path.getsize(path) / 1e6
    del net, pruned, variant_masks
    torch.cuda.empty_cache()
    variant_results[label] = row
    lines = [
        (f"Dense ({DENSE_EPOCHS} ep ft)", f"{row['dense']:.2f}%"),
        ("Pruned (no ft)", f"{row['pruned_no_ft']:.2f}%"),
        (f"Pruned ({FT_EPOCHS} ep ft)", f"{row['pruned_ft']:.2f}%"),
        ("Drop vs dense", f"{row['dense'] - row['pruned_ft']:.2f} pts"),
        ("Density (conv weights)", f"{row['density']:.1%} over {row['layers']} layers"),
        ("Masks valid, zeros kept", row["masks_ok"]),
        ("Minutes", f"{row['minutes']:.1f}"),
        ("Checkpoint", f"{path} ({row['checkpoint_mb']:.0f} MB)"),
    ]
    print("=" * 60)
    print(f"  {label}")
    for label_, value in lines:
        print(f"  {label_ + ':':<26}{value}")
    print("=" * 60)
    if DOWNLOAD_CHECKPOINTS and colab_files is not None:
        colab_files.download(path)

#@title ResNet-34 · IMAGENET1K_V1
run_variant("resnet34", "IMAGENET1K_V1")

#@title ResNet-50 · IMAGENET1K_V1
run_variant("resnet50", "IMAGENET1K_V1")

#@title ResNet-50 · IMAGENET1K_V2
run_variant("resnet50", "IMAGENET1K_V2")

#@title ResNet-101 · IMAGENET1K_V1
run_variant("resnet101", "IMAGENET1K_V1", batch=64)  # batch 64 to fit a T4's 15 GB (not measured)

#@title ResNet-101 · IMAGENET1K_V2
run_variant("resnet101", "IMAGENET1K_V2", batch=64)  # batch 64 to fit a T4's 15 GB (not measured)

#@title ResNet-152 · IMAGENET1K_V1
run_variant("resnet152", "IMAGENET1K_V1", batch=64)  # batch 64 to fit a T4's 15 GB (not measured)

#@title ResNet-152 · IMAGENET1K_V2
run_variant("resnet152", "IMAGENET1K_V2", batch=64)  # batch 64 to fit a T4's 15 GB (not measured)

Baseline: 2 of every 4 channels

The ResNet-50 V2 run again with one change: the mask keeps 2 of every 4 consecutive output channels instead of never keeping two neighbors, the baseline .docs/experiment/Zeckendorf.py compares against on ResNet-20. Both masks switch off whole channels at about 50% density, so the comparison isolates the pattern. If this run also ends far below its dense accuracy, the V2 drop is not specific to the φ pattern. This is not NVIDIA's hardware 2:4 sparsity, which keeps 2 of every 4 weights inside each row rather than whole channels. It needs a phi-prune with prune(pattern=...).

#@title ResNet-50 · IMAGENET1K_V2 · 2:4 channels
run_variant("resnet50", "IMAGENET1K_V2", pattern="2:4")

#@title Compare the variants
table = {}
if "ft_results" in globals():  # the walkthrough's ResNet-18 from sections 3-7 (not timed)
    table["resnet18 / DEFAULT"] = {
        "params_m": sum(p.numel() for p in model.parameters()) / 1e6,
        "dense": dense_acc, "pruned_ft": ft_results["best_val_acc"],
        "density": stats["density"], "layers": stats["pruned_layers"], "minutes": float("nan"),
        "masks_ok": report["_summary"]["all_masks_valid"] and report["_summary"]["all_zeros_enforced"],
    }
table.update(variant_results)
header = f"{'variant':<30}{'params':>8}{'dense':>9}{'pruned':>9}{'drop':>7}{'density':>9}{'layers':>8}"
print(header + f"{'min':>7}")
for name, r in table.items():
    line = (f"{name:<30}{r['params_m']:>7.1f}M{r['dense']:>8.2f}%{r['pruned_ft']:>8.2f}%"
            f"{r['dense'] - r['pruned_ft']:>7.2f}{r['density']:>9.1%}{r['layers']:>8}{r['minutes']:>7.1f}")
    print(line if r["masks_ok"] else line + "  masks FAILED")
with open("variant_results.json", "w") as f:
    json.dump(table, f, indent=1)
print("Saved variant_results.json")

Encode and check the checkpoints

The next cell reads every checkpoint the variant cells saved in checkpoints/, with no retraining. For each one it measures three things: the accuracy after Fibonacci-encoding every pruned layer (section 6 encodes one layer), the share of simulated single-bit flips the adjacency check catches, and the bits per weight of the self-delimiting Fibonacci stream. Each output channel gets its own scale and offset (PER_CHANNEL): with one range per layer, the few outlying weights of the ImageNet V2 recipes stretch it so far that encoding changes the network's predictions entirely (scripts/probe_pretrained_encoding.py); set PER_CHANNEL = False to reproduce the per-layer numbers. If the installed encoder rounds weights up instead of to the nearest level, the cell says so, and the encoded accuracy includes that extra error. The checkpoints stay on the VM only while the session lasts: in a new session, run the Setup, data and helpers cells, then upload the downloaded checkpoints into a checkpoints/ folder from the Files pane.

#@title Encode and check the saved checkpoints
import glob
import inspect
import json
import os
import random
import time

import numpy as np
from phi_prune.encoding import FibonacciEncoder
from phi_prune.integrity import simulate_corruption

N_DIGITS = 10          # 144 levels, as in section 6
FLIPS_PER_LAYER = 200  # simulated single-bit flips per pruned layer
PER_CHANNEL = True     # a scale and offset per output channel; False gives one per layer, as before 0.2.3

fib = FibonacciEncoder(n_digits=N_DIGITS)
if "axis" not in inspect.signature(fib.encode_tensor).parameters:
    raise RuntimeError("phi-prune 0.2.3+ is needed for per-channel encoding: rerun the install cell, "
                       "then restart the runtime")
probe = fib.encode_tensor(torch.tensor([0.0, 4.3, float(fib.max_value)]))[0]
if probe[1].item() != 4.0:  # on this 0..max_value range, 4.3 should snap to the nearest level, 4
    print("Note: this phi-prune's encoder rounds weights up to the next level instead of the\n"
          "nearest one, so the encoded accuracy below includes that extra error\n")
# stream bits for each grid value (encode_to_stream codes v + 1, so 0 has a codeword too)
code_bits = np.array([len(fib.encode_to_stream([v])) for v in range(fib.max_value + 1)])


def blank_model(arch):
    """The checkpoint's architecture with a 10-class fc and no downloaded weights."""
    net = torchvision.models.get_model(arch, weights=None)
    net.fc = nn.Linear(net.fc.in_features, 10)
    return net


def encode_checkpoint(path):
    """Accuracy before and after encoding every pruned layer, flips caught, stream bits per weight."""
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    net = blank_model(ckpt["metadata"]["arch"])
    net.load_state_dict(ckpt["state_dict"])
    net = net.to(device, memory_format=torch.channels_last)
    before = evaluate(net, test_loader)
    params = dict(net.named_parameters())
    bits = weights = caught = flips = 0
    with torch.no_grad():
        for name, mask in ckpt["masks"].items():
            param, mask = params[name], mask.to(device)
            encoded, scale, _, offset = fib.encode_tensor(param.data, mask=mask, return_offset=True,
                                                          axis=0 if PER_CHANNEL else None)
            param.copy_(encoded)
            levels = torch.from_numpy(fib.levels(param.data, scale, offset, mask))
            counts = torch.bincount(levels, minlength=fib.max_value + 1).numpy()
            bits += int((counts * code_bits).sum())
            weights += int(counts.sum())
            hit, n = simulate_corruption(param.data.cpu(), mask.cpu(), fib, scale, offset,
                                         n_flips=FLIPS_PER_LAYER)
            caught += hit
            flips += n
    after = evaluate(net, test_loader)
    del net
    torch.cuda.empty_cache()
    return {"pruned": before, "encoded": after,
            "flips_caught": caught / flips, "bits_per_weight": bits / weights}


random.seed(0)  # simulate_corruption draws from the random module
paths = sorted(glob.glob(os.path.join(CHECKPOINT_DIR, "*_pruned.pt")))
if not paths:
    print(f"No checkpoints in {CHECKPOINT_DIR}/: run a variant cell first, or upload saved ones there")
encoding_results = {}
print(f"{'checkpoint':<30}{'pruned':>9}{'encoded':>9}{'drop':>7}"
      f"{'flips caught':>14}{'bits/weight':>13}{'min':>6}")
for path in paths:
    start = time.perf_counter()
    r = encode_checkpoint(path)
    r["minutes"] = (time.perf_counter() - start) / 60
    name = os.path.basename(path)[:-len("_pruned.pt")]
    encoding_results[name] = r
    print(f"{name:<30}{r['pruned']:>8.2f}%{r['encoded']:>8.2f}%{r['pruned'] - r['encoded']:>7.2f}"
          f"{r['flips_caught']:>14.1%}{r['bits_per_weight']:>13.2f}{r['minutes']:>6.1f}")
with open("encoding_results.json", "w") as f:
    json.dump(encoding_results, f, indent=1)
print(f"Saved encoding_results.json ({fib.n_levels} levels; a fixed-width code needs "
      f"{np.log2(fib.n_levels):.2f} bits per weight)")