Add train_pdp.py

train_pdp.py

@@ -0,0 +1,297 @@
+"""
+PDP Training Script for CIFAR-10 with ResNet18
+Based on: PDP: Parameter-free Differentiable Pruning is All You Need (NeurIPS 2023)
+"""
+
+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 torchvision import transforms
+from datasets import load_dataset
+import numpy as np
+import argparse
+import json
+import os
+from tqdm import tqdm
+
+from pdp import PDPPruner
+
+
+# ---------------------------------------------------------------------------
+# CIFAR-10 adapted ResNet18
+# ---------------------------------------------------------------------------
+
+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 = 64
+        # First conv adapted for 32x32 CIFAR-10
+        self.conv1 = conv3x3(3, 64)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512 * 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 = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+def ResNet18(num_classes=10):
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes)
+
+
+# ---------------------------------------------------------------------------
+# Data loading
+# ---------------------------------------------------------------------------
+
+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
+
+
+# ---------------------------------------------------------------------------
+# Training & evaluation helpers
+# ---------------------------------------------------------------------------
+
+def train_epoch(model, loader, optimizer, criterion, device, pruner=None, epoch=None):
+    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()
+
+        if pruner is not None and epoch is not None:
+            pruner.step(epoch)
+
+        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
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="PDP Training on CIFAR-10")
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--lr", type=float, default=0.1)
+    parser.add_argument("--momentum", type=float, default=0.9)
+    parser.add_argument("--weight_decay", type=float, default=5e-4)
+    parser.add_argument("--target_sparsity", type=float, default=0.85)
+    parser.add_argument("--s", type=int, default=16, help="Warmup epochs before pruning starts")
+    parser.add_argument("--epsilon", type=float, default=0.015, help="Gradual pruning rate per epoch")
+    parser.add_argument("--tau", type=float, default=1e-4, help="PDP temperature")
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--save_dir", type=str, default="./checkpoints")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)
+    if args.device == "cuda":
+        torch.cuda.manual_seed(args.seed)
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    # Data
+    train_loader, test_loader = get_cifar10_loaders(args.batch_size, args.num_workers)
+    print(f"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}")
+
+    # Model
+    model = ResNet18(num_classes=10).to(device)
+    print(f"Model params: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Optimizer & scheduler
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
+                          weight_decay=args.weight_decay)
+    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60, 90], gamma=0.1)
+
+    # PDP Pruner
+    pruner = PDPPruner(
+        model=model,
+        target_sparsity=args.target_sparsity,
+        s=args.s,
+        epsilon=args.epsilon,
+        tau=args.tau,
+    )
+    pruner.attach()
+
+    # Training loop
+    history = []
+    best_acc = 0.0
+
+    for epoch in range(args.epochs):
+        train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device, pruner=pruner, epoch=epoch)
+        val_loss, val_acc = evaluate(model, test_loader, criterion, device)
+        scheduler.step()
+
+        current_sparsity = pruner.get_sparsity()
+        effective = pruner.current_effective_sparsity
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs} | "
+              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | "
+              f"Val Loss: {val_loss:.4f} Acc: {val_acc:.2f}% | "
+              f"Sparsity: {current_sparsity:.4f} (eff: {effective:.4f}) | "
+              f"LR: {optimizer.param_groups[0]['lr']:.4f}")
+
+        history.append({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "train_acc": train_acc,
+            "val_loss": val_loss,
+            "val_acc": val_acc,
+            "sparsity": current_sparsity,
+            "effective_sparsity": effective,
+            "lr": optimizer.param_groups[0]["lr"],
+        })
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            ckpt_path = os.path.join(args.save_dir, "best_model.pt")
+            torch.save({
+                "epoch": epoch + 1,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "pruner_state_dict": pruner.state_dict(),
+                "val_acc": val_acc,
+            }, ckpt_path)
+
+    # Final hard prune and evaluation
+    print("\n--- Final Hard Pruning ---")
+    pruner.hard_prune()
+    final_sparsity = pruner.get_sparsity()
+    final_val_loss, final_val_acc = evaluate(model, test_loader, criterion, device)
+    print(f"After hard prune: Sparsity={final_sparsity:.4f}, Val Acc={final_val_acc:.2f}%")
+
+    # Save final model
+    final_path = os.path.join(args.save_dir, "final_model.pt")
+    torch.save({
+        "model_state_dict": model.state_dict(),
+        "pruner_state_dict": pruner.state_dict(),
+        "final_sparsity": final_sparsity,
+        "final_val_acc": final_val_acc,
+    }, final_path)
+
+    # Save history
+    with open(os.path.join(args.save_dir, "history.json"), "w") as f:
+        json.dump(history, f, indent=2)
+
+    print(f"\nBest validation accuracy: {best_acc:.2f}%")
+    print(f"Final pruned model saved to {final_path}")
+
+
+if __name__ == "__main__":
+    main()