Add training script

train.py

@@ -0,0 +1,278 @@
+"""
+DKM Training Pipeline
+
+Implements the training protocol from Section 4 of the paper:
+- Start from pre-trained model
+- Insert DKM layers for weight clustering
+- Fine-tune with SGD (momentum=0.9, lr=0.008)
+- No loss function or architecture modifications
+
+This demo uses CIFAR-10 with a ResNet model to demonstrate the full pipeline.
+For ImageNet reproduction, scale up to the full dataset and use 8xV100 GPUs.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+import time
+import argparse
+import json
+import os
+
+from dkm import DKMCompressor, compress_model
+from dkm.utils import print_compression_summary, count_unique_weights
+
+
+def get_cifar10_loaders(batch_size=128, num_workers=2):
+    """Get CIFAR-10 train/test data loaders with standard augmentation."""
+    transform_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+    ])
+    
+    transform_test = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+    ])
+    
+    trainset = torchvision.datasets.CIFAR10(
+        root='./data', train=True, download=True, transform=transform_train
+    )
+    trainloader = DataLoader(
+        trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers
+    )
+    
+    testset = torchvision.datasets.CIFAR10(
+        root='./data', train=False, download=True, transform=transform_test
+    )
+    testloader = DataLoader(
+        testset, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
+    
+    return trainloader, testloader
+
+
+def evaluate(model, dataloader, device, criterion=None):
+    """Evaluate model accuracy and optional loss."""
+    model.eval()
+    correct = 0
+    total = 0
+    total_loss = 0.0
+    n_batches = 0
+    
+    with torch.no_grad():
+        for images, labels in dataloader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            
+            if criterion:
+                loss = criterion(outputs, labels)
+                total_loss += loss.item()
+                n_batches += 1
+            
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+    
+    accuracy = 100.0 * correct / total
+    avg_loss = total_loss / max(n_batches, 1)
+    return accuracy, avg_loss
+
+
+def train_dkm(
+    model,
+    trainloader,
+    testloader,
+    device,
+    bits=2,
+    dim=1,
+    tau=2e-5,
+    epochs=50,
+    lr=0.008,
+    momentum=0.9,
+    weight_decay=0.0,
+    skip_first_last=True,
+    conv_config=None,
+    fc_config=None,
+):
+    """
+    Train a model with DKM compression.
+    
+    Follows the paper's protocol:
+    - SGD optimizer with momentum=0.9
+    - Fixed learning rate (no per-layer tuning)
+    - Original loss function (CrossEntropyLoss)
+    - DKM inserted into forward pass
+    """
+    print(f"\\n{'='*70}")
+    print(f"DKM Compression Training")
+    print(f"{'='*70}")
+    print(f"Bits: {bits}, Dim: {dim}, Tau: {tau}")
+    print(f"Epochs: {epochs}, LR: {lr}, Momentum: {momentum}")
+    print(f"Skip first/last: {skip_first_last}")
+    if conv_config:
+        print(f"Conv config: {conv_config}")
+    if fc_config:
+        print(f"FC config: {fc_config}")
+    
+    # Evaluate baseline (pre-trained) accuracy first
+    model = model.to(device)
+    criterion = nn.CrossEntropyLoss()
+    
+    print("\\nEvaluating baseline (pre-trained) model...")
+    baseline_acc, baseline_loss = evaluate(model, testloader, device, criterion)
+    print(f"Baseline accuracy: {baseline_acc:.2f}%, Loss: {baseline_loss:.4f}")
+    
+    # Wrap model with DKM compression
+    compressor = compress_model(
+        model=model,
+        bits=bits,
+        dim=dim,
+        tau=tau,
+        conv_config=conv_config,
+        fc_config=fc_config,
+        skip_first_last=skip_first_last,
+    )
+    compressor = compressor.to(device)
+    
+    # Print compression info
+    info = compressor.get_compression_info()
+    print_compression_summary(info)
+    
+    # SGD optimizer (paper: momentum=0.9, lr=0.008)
+    optimizer = optim.SGD(
+        compressor.parameters(), 
+        lr=lr, 
+        momentum=momentum,
+        weight_decay=weight_decay,
+    )
+    
+    # Training loop
+    best_acc = 0.0
+    history = []
+    
+    for epoch in range(epochs):
+        compressor.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+        epoch_start = time.time()
+        
+        for batch_idx, (images, labels) in enumerate(trainloader):
+            images, labels = images.to(device), labels.to(device)
+            
+            optimizer.zero_grad()
+            outputs = compressor(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            
+            running_loss += loss.item()
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+        
+        train_acc = 100.0 * correct / total
+        train_loss = running_loss / len(trainloader)
+        epoch_time = time.time() - epoch_start
+        
+        # Evaluate (with hard assignment for accurate eval)
+        test_acc, test_loss = evaluate(compressor, testloader, device, criterion)
+        
+        if test_acc > best_acc:
+            best_acc = test_acc
+        
+        history.append({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "train_acc": train_acc,
+            "test_loss": test_loss,
+            "test_acc": test_acc,
+            "best_acc": best_acc,
+            "epoch_time": epoch_time,
+        })
+        
+        print(
+            f"Epoch [{epoch+1}/{epochs}] "
+            f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | "
+            f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}% | "
+            f"Best: {best_acc:.2f}% | Time: {epoch_time:.1f}s"
+        )
+    
+    # Final: snap weights to centroids
+    print("\\nSnapping weights to nearest centroids...")
+    compressor.snap_weights()
+    
+    # Verify unique weights
+    unique_counts = count_unique_weights(model)
+    print("\\nUnique weight values per layer after compression:")
+    for name, count in unique_counts.items():
+        print(f"  {name}: {count} unique values")
+    
+    final_acc, final_loss = evaluate(compressor, testloader, device, criterion)
+    print(f"\\nFinal (snapped) accuracy: {final_acc:.2f}%")
+    print(f"Accuracy drop from baseline: {baseline_acc - final_acc:.2f}%")
+    print(f"Best training accuracy: {best_acc:.2f}%")
+    
+    return compressor, history, info
+
+
+def main():
+    parser = argparse.ArgumentParser(description="DKM Compression Training")
+    parser.add_argument("--bits", type=int, default=2, help="Number of bits for clustering")
+    parser.add_argument("--dim", type=int, default=1, help="Clustering dimension")
+    parser.add_argument("--tau", type=float, default=2e-5, help="Temperature parameter")
+    parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs")
+    parser.add_argument("--lr", type=float, default=0.008, help="Learning rate")
+    parser.add_argument("--batch-size", type=int, default=128, help="Batch size")
+    parser.add_argument("--device", type=str, default="auto", help="Device (auto/cpu/cuda)")
+    parser.add_argument("--save-path", type=str, default="dkm_compressed.pt")
+    args = parser.parse_args()
+    
+    # Device
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+    print(f"Using device: {device}")
+    
+    # Load pre-trained ResNet18 
+    print("Loading pre-trained ResNet18...")
+    model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
+    # Adapt for CIFAR-10 (10 classes instead of 1000)
+    model.fc = nn.Linear(model.fc.in_features, 10)
+    
+    # Get data
+    trainloader, testloader = get_cifar10_loaders(batch_size=args.batch_size)
+    
+    # Train with DKM
+    compressor, history, info = train_dkm(
+        model=model,
+        trainloader=trainloader,
+        testloader=testloader,
+        device=device,
+        bits=args.bits,
+        dim=args.dim,
+        tau=args.tau,
+        epochs=args.epochs,
+        lr=args.lr,
+    )
+    
+    # Save
+    export = compressor.export_compressed()
+    torch.save(export, args.save_path)
+    print(f"\\nCompressed model saved to {args.save_path}")
+    
+    # Save history
+    with open("training_history.json", "w") as f:
+        json.dump(history, f, indent=2)
+    print("Training history saved to training_history.json")
+
+
+if __name__ == "__main__":
+    main()