train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import numpy as np
from torch.optim.lr_scheduler import MultiStepLR

# Set random seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed(42)


def ResNet56(model_name="resnet56", num_classes=100):
    model = torch.hub.load("chenyaofo/pytorch-cifar-models",
                           f"cifar{num_classes}_{model_name}", pretrained=False)
    return model


# Data augmentation and preprocessing
def get_dataloaders(batch_size=128):
    # Training transforms with strong augmentation
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.AutoAugment(transforms.AutoAugmentPolicy.CIFAR10),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
    ])

    # Test transforms
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
    ])

    trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)

    testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
    testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=4, pin_memory=True)

    return trainloader, testloader


# Training function
def train_epoch(model, trainloader, criterion, optimizer, device):
    model.train()
    train_loss = 0
    correct = 0
    total = 0

    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    return train_loss / len(trainloader), 100. * correct / total


# Testing function
def test(model, testloader, criterion, device):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    return test_loss / len(testloader), 100. * correct / total


# Main training loop
def main():
    # Hyperparameters (optimized for best accuracy)
    batch_size = 128
    epochs = 200
    initial_lr = 0.1
    momentum = 0.9
    weight_decay = 5e-4
    milestones = [60, 120, 160]  # Learning rate decay milestones
    gamma = 0.2  # Learning rate decay factor

    model_name = "repvgg_a2"  # Model name for saving

    # Setup
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # Load data
    print("Loading CIFAR-100 dataset...")
    trainloader, testloader = get_dataloaders(batch_size)

    # Create model
    print(f"Creating {model_name} model...")
    model = ResNet56(model_name=model_name, num_classes=100).to(device)
    
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)  # Label smoothing for better generalization
    optimizer = optim.SGD(model.parameters(), lr=initial_lr, momentum=momentum, weight_decay=weight_decay)
    scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=gamma)

    # Training loop
    best_acc = 0
    print(f"\nStarting training for {epochs} epochs...")
    print("-" * 80)

    for epoch in range(epochs):
        train_loss, train_acc = train_epoch(model, trainloader, criterion, optimizer, device)
        test_loss, test_acc = test(model, testloader, criterion, device)
        scheduler.step()

        # Save best model
        if test_acc > best_acc:
            print(f"Saving model (accuracy improved: {best_acc:.2f}% -> {test_acc:.2f}%)")
            best_acc = test_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_acc': best_acc,
            }, f'/home/microway/Selective Classification/resnet_train/{model_name}_cifar100_best.pth')

        # Print progress
        if (epoch + 1) % 10 == 0 or epoch == 0:
            print(f"Epoch: {epoch+1:3d} | LR: {scheduler.get_last_lr()[0]:.6f} | "
                  f"Train Loss: {train_loss:.3f} | Train Acc: {train_acc:.2f}% | "
                  f"Test Loss: {test_loss:.3f} | Test Acc: {test_acc:.2f}% | "
                  f"Best: {best_acc:.2f}%")

    print("-" * 80)
    print(f"\nTraining completed! Best test accuracy: {best_acc:.2f}%")
    print(f"Best model saved to: resnet56_cifar100_best.pth")


if __name__ == '__main__':
    main()