app.py

# Libraries
import time
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms


# ============================================================================
# PAPER IMPLEMENTATION: Perceptron-based Pooling Operations
# ============================================================================

class PerceptronPool2d(nn.Module):
    """
    Single perceptron as pooling operation (Paper's core contribution)
    - No activation function (paper found this works better - see Table 2)
    - Parameters: W*H weights + 1 bias
    """
    def __init__(self, kernel_size=2, stride=2):
        super(PerceptronPool2d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        
        # Initializing weights similar to average pooling (0.25 for 2x2)
        # but with small random variations as per paper's initialization strategy
        init_value = 1.0 / (kernel_size * kernel_size)
        self.weight = nn.Parameter(
            torch.randn(kernel_size * kernel_size) * 0.1 + init_value
        )
        self.bias = nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        # x shape: [batch, channels, height, width]
        batch, channels, h, w = x.shape
        
        # Calculating output dimensions
        out_h = (h - self.kernel_size) // self.stride + 1
        out_w = (w - self.kernel_size) // self.stride + 1
        
        # Extracting sliding windows using unfold
        # Output: [batch, channels * kernel_size^2, num_patches]
        patches = nn.functional.unfold(
            x, 
            kernel_size=self.kernel_size, 
            stride=self.stride
        )
        
        # Reshape: [batch, channels, kernel_size^2, num_patches]
        patches = patches.reshape(batch, channels, self.kernel_size**2, -1)
        
        # Appling perceptron: weighted sum + bias (NO activation function)
        output = torch.einsum('bckp,k->bcp', patches, self.weight) + self.bias
        
        # Reshaping back to spatial dimensions
        output = output.reshape(batch, channels, out_h, out_w)
        
        return output


class NeuralNetPool2d(nn.Module):
    """
    Multi-layer neural network as pooling (NN-4-1 from Table 3)
    - Hidden layer: 4 perceptrons (no activation)
    - Output layer: 1 perceptron (no activation)
    - Total: 50 additional parameters
    """
    def __init__(self, kernel_size=2, stride=2, hidden_neurons=4):
        super(NeuralNetPool2d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        input_size = kernel_size * kernel_size
        
        # First layer: 4 perceptrons
        init_value = 1.0 / input_size
        self.layer1_weights = nn.Parameter(
            torch.randn(hidden_neurons, input_size) * 0.1 + init_value
        )
        self.layer1_bias = nn.Parameter(torch.zeros(hidden_neurons))
        
        # Output layer: 1 perceptron
        self.layer2_weights = nn.Parameter(
            torch.randn(1, hidden_neurons) * 0.1 + (1.0 / hidden_neurons)
        )
        self.layer2_bias = nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        batch, channels, h, w = x.shape
        out_h = (h - self.kernel_size) // self.stride + 1
        out_w = (w - self.kernel_size) // self.stride + 1
        
        # Extracting patches
        patches = nn.functional.unfold(
            x, kernel_size=self.kernel_size, stride=self.stride
        )
        patches = patches.reshape(batch, channels, self.kernel_size**2, -1)
        
        # Hidden layer (no activation function)
        hidden = torch.einsum('bckp,hk->bchp', patches, self.layer1_weights)
        hidden = hidden + self.layer1_bias.view(1, 1, -1, 1)
        
        # Output layer (no activation function)
        output = torch.einsum('bchp,oh->bcop', hidden, self.layer2_weights)
        output = output + self.layer2_bias.view(1, 1, -1, 1)
        output = output.squeeze(2).reshape(batch, channels, out_h, out_w)
        
        return output


# ============================================================================
# Network Architecture
# ============================================================================

class Net(nn.Module):
    """
    Network with perceptron-based pooling
    Choose pool_type: 'perceptron', 'nn-4-1', or 'max' (baseline)
    """
    def __init__(self, pool_type='perceptron'):
        super(Net, self).__init__()
        
        # Convolutional layers (unchanged)
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        
        # Pooling layers - THIS IS THE KEY CHANGE COMPARED TO THE LAST CODE I WROTE
        if pool_type == 'perceptron':
            print("Using Perceptron Pooling (10 params per layer)")
            self.pool1 = PerceptronPool2d(kernel_size=2, stride=2)
            self.pool2 = PerceptronPool2d(kernel_size=2, stride=2)
            self.pool3 = PerceptronPool2d(kernel_size=2, stride=2)
        elif pool_type == 'nn-4-1':
            print("Using NN-4-1 Pooling (50 params per layer)")
            self.pool1 = NeuralNetPool2d(kernel_size=2, stride=2, hidden_neurons=4)
            self.pool2 = NeuralNetPool2d(kernel_size=2, stride=2, hidden_neurons=4)
            self.pool3 = NeuralNetPool2d(kernel_size=2, stride=2, hidden_neurons=4)
        else:  # max pooling baseline
            print("Using Max Pooling (baseline)")
            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Fully connected layer
        self.fc1 = nn.Linear(128 * 4 * 4, 100)
    
    def forward(self, x):
        x = self.bn1(torch.relu(self.conv1(x)))
        x = self.pool1(x)
        x = self.bn2(torch.relu(self.conv2(x)))
        x = self.pool2(x)
        x = self.bn3(torch.relu(self.conv3(x)))
        x = self.pool3(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        return x


# ============================================================================
# Improved Initialization and Training
# ============================================================================

def weights_init(m):
    """Initializing weights for conv and fc layers"""
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        torch.nn.init.kaiming_normal_(m.weight)
        if m.bias is not None:
            torch.nn.init.constant_(m.bias, 0)


def evaluate(model, dataloader, device):
    """Evaluating model accuracy on validation/test set"""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    model.train()
    return 100 * correct / total


# ============================================================================
# Main Training Loop (Following Paper's CIFAR-100 Setup)
# ============================================================================

def main():
    start_time = time.perf_counter()
    
    # ========================================================================
    # Data Setup (Following paper specifications - page 5)
    # ========================================================================
    
    # Training transforms: normalize to zero mean and unit std + data augmentation
    transform_train = transforms.Compose([
        transforms.Pad(4, fill=0),  # Paper: 4 pixels zero padding on each side
        transforms.RandomCrop(32),   # Paper: crop 32x32 from 40x40
        transforms.ToTensor(),
        # Paper specifies: normalize to zero mean and one standard deviation
        # Using CIFAR-100 statistics
        transforms.Normalize(
            mean=[0.5071, 0.4867, 0.4408],
            std=[0.2675, 0.2565, 0.2761]
        )
    ])
    
    # Test transforms: only normalize (no augmentation)
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.5071, 0.4867, 0.4408],
            std=[0.2675, 0.2565, 0.2761]
        )
    ])
    
    # Loading datasets
    trainset = torchvision.datasets.CIFAR100(
        root='./data', train=True, download=True, transform=transform_train
    )
    testset = torchvision.datasets.CIFAR100(
        root='./data', train=False, download=True, transform=transform_test
    )
    
    # Data loaders
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=100, shuffle=True, num_workers=2, 
        pin_memory=torch.cuda.is_available()
    )
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=100, shuffle=False, num_workers=2,
        pin_memory=torch.cuda.is_available()
    )
    
    # ========================================================================
    # Model Setup
    # ========================================================================
    
    device = torch.device(
        "cuda" if torch.cuda.is_available() 
        else "mps" if torch.backends.mps.is_available() 
        else "cpu"
    )
    print(f"Using device: {device}")
    
    # Pooling
    net = Net(pool_type='perceptron').to(device)
    net.apply(weights_init)
    
    # ========================================================================
    # Training Setup (Following paper specifications - page 5)
    # ========================================================================
    
    criterion = nn.CrossEntropyLoss()
    
    # Paper specifies for CIFAR-100:
    # - Optimizer: SGD with momentum (0.9)
    # - Learning rate: 0.1
    # - Weight decay: 5e-4
    # - LR schedule: reduce by 0.1 at epochs 80 and 120
    # - Total epochs: 160
    
    # Separating pooling parameters for reduced learning rate
    pooling_params = []
    other_params = []
    
    for name, param in net.named_parameters():
        if 'pool' in name:
            pooling_params.append(param)
        else:
            other_params.append(param)
    
    # Paper: reduce learning rate by factor 10^-1 for pooling layers
    optimizer = torch.optim.SGD([
        {'params': other_params},
        {'params': pooling_params, 'lr': 0.01, 'weight_decay': 0}  # Factor 0.1
    ], lr=0.1, momentum=0.9, weight_decay=5e-4)
    
    # Paper schedule: reduce LR at epochs 80 and 120
    total_epochs = 160
    
    print(f"\nTraining for {total_epochs} epochs")
    print(f"Pooling parameters: {len(pooling_params)}")
    print(f"Other parameters: {len(other_params)}")
    print("-" * 60)
    
    # ========================================================================
    # Training Loop
    # ========================================================================
    
    best_acc = 0.0
    
    for epoch in range(total_epochs):
        # Adjusting learning rate at specified epochs
        if epoch == 80 or epoch == 120:
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 0.1
            print(f"Learning rate reduced at epoch {epoch + 1}")
        
        net.train()
        running_loss = 0.0
        total_loss = 0.0  # Track full epoch loss (fixed bug)
        
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            
            # Zero gradients
            optimizer.zero_grad()
            
            # Forward pass
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            
            # Backward pass
            loss.backward()
            optimizer.step()
            
            # Track loss
            running_loss += loss.item()
            total_loss += loss.item()
            
            # Print every 200 mini-batches
            if i % 200 == 199:
                avg_loss = running_loss / 200
                print(f'[{epoch + 1:3d}, {i + 1:5d}] loss: {avg_loss:.3f}')
                running_loss = 0.0
            
            # Safety check
            if torch.isnan(loss) or torch.isinf(loss):
                print(f"NaN or Inf detected in loss at epoch {epoch}, batch {i}")
                return
        
        # Calculating epoch statistics
        epoch_loss = total_loss / len(trainloader)
        
        # Evaluating on test set every 5 epochs
        if (epoch + 1) % 5 == 0 or epoch == 0:
            test_acc = evaluate(net, testloader, device)
            print(f'Epoch {epoch + 1}/{total_epochs} | '
                  f'Loss: {epoch_loss:.3f} | '
                  f'Test Acc: {test_acc:.2f}%')
            
            if test_acc > best_acc:
                best_acc = test_acc
                # Optionally save best model
                # torch.save(net.state_dict(), 'best_model.pth')
        else:
            print(f'Epoch {epoch + 1}/{total_epochs} | Loss: {epoch_loss:.3f}')
    
    # ========================================================================
    # Final Evaluation
    # ========================================================================
    
    final_acc = evaluate(net, testloader, device)
    
    end_time = time.perf_counter()
    elapsed_time = end_time - start_time
    
    print("\n" + "=" * 60)
    print(f'Training Completed!')
    print(f'Time taken: {elapsed_time:.2f} seconds ({elapsed_time/60:.2f} minutes)')
    print(f'Best test accuracy: {best_acc:.2f}%')
    print(f'Final test accuracy: {final_acc:.2f}%')
    print("=" * 60)
    
    # Calculation of additional parameters from pooling
    pooling_param_count = sum(p.numel() for p in pooling_params)
    print(f'\nAdditional parameters from pooling: {pooling_param_count}')
    
    return net, final_acc


if __name__ == '__main__':
    main()