metapruning/inference.py

"""
MetaPruning Inference Pipeline.

1. Load trained metanetwork
2. Take any target model
3. Convert -> metanetwork feedforward -> transform back
4. Finetune the transformed model
5. Prune using magnitude-based criterion
6. Evaluate

Paper: "Meta Pruning via Graph Metanetworks" (arXiv:2506.12041)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from datasets import load_dataset
from torchvision import transforms
import argparse
import os

from graph import resnet_to_graph, create_transformed_model, Graph
from gnn import Metanetwork


# ---------------------------------------------------------------------------
# Data loading (same as training script)
# ---------------------------------------------------------------------------

def get_cifar10_loaders(batch_size=128, num_workers=4):
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])

    ds_train = load_dataset("uoft-cs/cifar10", split="train")
    ds_test = load_dataset("uoft-cs/cifar10", split="test")

    def map_train(examples):
        images = [transform_train(img.convert("RGB")) for img in examples["img"]]
        return {"pixel_values": images, "labels": examples["label"]}

    def map_test(examples):
        images = [transform_test(img.convert("RGB")) for img in examples["img"]]
        return {"pixel_values": images, "labels": examples["label"]}

    ds_train = ds_train.map(map_train, batched=True, remove_columns=["img", "label"])
    ds_test = ds_test.map(map_test, batched=True, remove_columns=["img", "label"])

    ds_train.set_format(type="torch", columns=["pixel_values", "labels"])
    ds_test.set_format(type="torch", columns=["pixel_values", "labels"])

    train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)
    test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

    return train_loader, test_loader


# ---------------------------------------------------------------------------
# Model definitions (from training script)
# ---------------------------------------------------------------------------

def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self, in_planes, planes, stride=1):
        super().__init__()
        self.conv1 = conv3x3(in_planes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1,
                          stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion * planes)
            )
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10):
        super().__init__()
        self.in_planes = 16
        self.conv1 = conv3x3(3, 16)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
        self.linear = nn.Linear(64 * block.expansion, num_classes)
    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for s in strides:
            layers.append(block(self.in_planes, planes, s))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = F.avg_pool2d(out, out.size()[3])
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


def ResNet56(num_classes=10):
    return ResNet(BasicBlock, [9, 9, 9], num_classes=num_classes)


def ResNet110(num_classes=10):
    return ResNet(BasicBlock, [18, 18, 18], num_classes=num_classes)


# ---------------------------------------------------------------------------
# Training helpers
# ---------------------------------------------------------------------------

def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0
    for batch in loader:
        inputs, targets = batch["pixel_values"].to(device), batch["labels"].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    return total_loss / total, 100.0 * correct / total


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    for batch in loader:
        inputs, targets = batch["pixel_values"].to(device), batch["labels"].to(device)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        total_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
    return total_loss / total, 100.0 * correct / total


# ---------------------------------------------------------------------------
# Simple magnitude-based structured pruning
# ---------------------------------------------------------------------------

def prune_model(model, target_sparsity=0.5):
    """
    Simple channel pruning based on L2 norm of conv filter weights.
    For a proper implementation, use torch-pruning with DepGraph.
    """
    # Compute L2 norm per output channel for each conv layer
    channel_norms = {}
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            norms = module.weight.data.view(module.out_channels, -1).norm(dim=1)
            channel_norms[name] = norms
    
    # For simplicity, just compute a global threshold across all channels
    all_norms = torch.cat([n for n in channel_norms.values()])
    threshold_idx = int(target_sparsity * all_norms.numel())
    threshold = torch.sort(all_norms)[0][threshold_idx].item()
    
    # Prune channels below threshold
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            norms = channel_norms[name]
            keep_mask = norms > threshold
            num_keep = keep_mask.sum().item()
            if num_keep < module.out_channels:
                # Simple: just zero out pruned channels
                for ch in range(module.out_channels):
                    if not keep_mask[ch]:
                        module.weight.data[ch] = 0
                        if module.bias is not None:
                            module.bias.data[ch] = 0
    
    print(f"[Prune] Applied simple magnitude pruning (target={target_sparsity:.2f})")


def compute_model_sparsity(model):
    total = 0
    zeros = 0
    for p in model.parameters():
        total += p.numel()
        zeros += (p.data == 0).sum().item()
    return zeros / total


# ---------------------------------------------------------------------------
# Main inference pipeline
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="MetaPruning Inference Pipeline")
    parser.add_argument("--metanetwork_path", type=str, required=True,
                        help="Path to trained metanetwork checkpoint")
    parser.add_argument("--target_model", type=str, default="resnet56",
                        choices=["resnet56", "resnet110"])
    parser.add_argument("--finetune_epochs", type=int, default=100,
                        help="Finetune epochs after metanetwork (paper uses 100-200)")
    parser.add_argument("--prune_sparsity", type=float, default=0.5,
                        help="Target pruning sparsity")
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--lr", type=float, default=0.01)
    parser.add_argument("--momentum", type=float, default=0.9)
    parser.add_argument("--weight_decay", type=float, default=5e-4)
    parser.add_argument("--milestones", type=int, nargs="+", default=[60, 90])
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device(args.device)
    print(f"Using device: {device}")

    # Load data
    train_loader, test_loader = get_cifar10_loaders(args.batch_size, args.num_workers)

    # Load target model
    if args.target_model == "resnet56":
        model = ResNet56(num_classes=10).to(device)
    else:
        model = ResNet110(num_classes=10).to(device)
    print(f"Loaded target model: {args.target_model}")

    # Load metanetwork
    ckpt = torch.load(args.metanetwork_path, map_location=device)
    config = ckpt["config"]

    metanetwork = Metanetwork(
        node_in_dim=config["node_in_dim"],
        edge_in_dim=config["edge_in_dim"],
        node_out_dim=config["node_out_dim"],
        edge_out_dim=config["edge_out_dim"],
        hidden_dim=config["hidden_dim"],
        num_layers=config["num_layers"],
        alpha=config["alpha"],
        beta=config["beta"],
    ).to(device)
    metanetwork.load_state_dict(ckpt["metanetwork_state_dict"])
    metanetwork.eval()
    print(f"Loaded metanetwork (hidden_dim={config['hidden_dim']}, layers={config['num_layers']})")

    # Baseline: evaluate untransformed model
    criterion = nn.CrossEntropyLoss()
    _, base_acc = evaluate(model, test_loader, criterion, device)
    print(f"\nBaseline model accuracy (before metanetwork): {base_acc:.2f}%")

    # Step 1: Convert to graph
    print("\n[Step 1] Converting model to graph...")
    graph = resnet_to_graph(model, max_kernel_size=3)
    print(f"  Nodes: {graph.node_features.size(0)}, Edges: {graph.edge_features.size(0)}")

    # Step 2: Feed through metanetwork
    print("[Step 2] Metanetwork feedforward...")
    with torch.no_grad():
        graph.node_features = graph.node_features.to(device)
        graph.edge_features = graph.edge_features.to(device)
        graph.edge_index = graph.edge_index.to(device)

        gnn_output = metanetwork(
            graph.node_features,
            graph.edge_index,
            graph.edge_features,
        )

    # Step 3: Convert back to transformed model
    print("[Step 3] Converting transformed graph back to model...")
    transformed_model = create_transformed_model(graph, gnn_output, model).to(device)

    # Evaluate after metanetwork (before finetuning)
    _, meta_acc = evaluate(transformed_model, test_loader, criterion, device)
    print(f"  Accuracy after metanetwork (before finetune): {meta_acc:.2f}%")

    # Step 4: Finetune transformed model
    print(f"\n[Step 4] Finetuning for {args.finetune_epochs} epochs...")
    optimizer = optim.SGD(transformed_model.parameters(), lr=args.lr,
                          momentum=args.momentum, weight_decay=args.weight_decay)
    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.1)

    best_acc = 0.0
    for epoch in range(args.finetune_epochs):
        train_loss, train_acc = train_epoch(transformed_model, train_loader, optimizer, criterion, device)
        val_loss, val_acc = evaluate(transformed_model, test_loader, criterion, device)
        scheduler.step()

        if val_acc > best_acc:
            best_acc = val_acc

        if (epoch + 1) % 20 == 0:
            print(f"  Epoch {epoch+1:3d}: train_acc={train_acc:.2f}%, val_acc={val_acc:.2f}%")

    print(f"  Best finetuned accuracy: {best_acc:.2f}%")

    # Step 5: Prune
    print(f"\n[Step 5] Pruning (target sparsity={args.prune_sparsity:.2f})...")
    prune_model(transformed_model, target_sparsity=args.prune_sparsity)
    sparsity = compute_model_sparsity(transformed_model)
    print(f"  Actual model sparsity: {sparsity:.4f}")

    # Evaluate pruned model
    _, pruned_acc = evaluate(transformed_model, test_loader, criterion, device)
    print(f"  Accuracy after pruning: {pruned_acc:.2f}%")

    # Summary
    print("\n" + "=" * 50)
    print("SUMMARY")
    print("=" * 50)
    print(f"Baseline accuracy:      {base_acc:.2f}%")
    print(f"After metanetwork:      {meta_acc:.2f}%")
    print(f"After finetuning:       {best_acc:.2f}%")
    print(f"After pruning:          {pruned_acc:.2f}%")
    print(f"Sparsity:               {sparsity:.4f}")
    print(f"Accuracy drop:          {base_acc - pruned_acc:.2f}%")
    print("=" * 50)


if __name__ == "__main__":
    main()