Untitled

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.cuda.amp import GradScaler, autocast

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. Model Setup ---------------------------------------------------------------
model = models.inception_v3(pretrained=True)
model.aux_logits = True  # Keep auxiliary classifier

# Unfreeze all layers
for param in model.parameters():
    param.requires_grad = True

# Replace final layers (both main and auxiliary)
num_classes = 2
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
    nn.Dropout(0.5),
    nn.Linear(num_ftrs, num_classes)
)

if model.AuxLogits is not None:
    num_aux_ftrs = model.AuxLogits.fc.in_features
    model.AuxLogits.fc = nn.Linear(num_aux_aux_ftrs, num_classes)

model = model.to(device)

# 2. Optimizer & Scheduler -----------------------------------------------------
optimizer = optim.AdamW(
    model.parameters(),
    lr=1e-4,            # Lower initial LR for fine-tuning
    weight_decay=1e-2   # Regularization
)

scheduler = ReduceLROnPlateau(
    optimizer,
    mode='max',         # Monitor validation accuracy
    factor=0.5,         # Reduce LR by 50%
    patience=2,         # Wait 2 epochs w/o improvement
    verbose=True
)

# 3. Data Augmentation ---------------------------------------------------------
train_transforms = transforms.Compose([
    transforms.Resize(342),          # For Inception-specific sizing
    transforms.RandomResizedCrop(299),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

val_transforms = transforms.Compose([
    transforms.Resize(342),
    transforms.CenterCrop(299),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 4. Training Loop with Best Practices -----------------------------------------
criterion = nn.CrossEntropyLoss()
scaler = GradScaler()  # For mixed precision training
best_acc = 0.0

for epoch in range(50):  # Maximum 50 epochs
    # Training Phase
    model.train()
    running_loss = 0.0
    
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        optimizer.zero_grad()
        
        # Mixed precision forward
        with autocast():
            outputs, aux_outputs = model(inputs)
            loss1 = criterion(outputs, labels)
            loss2 = criterion(aux_outputs, labels)
            loss = loss1 + 0.4 * loss2  # Aux loss weight
            
        # Backward pass with scaling
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        
        running_loss += loss.item() * inputs.size(0)
    
    # Validation Phase
    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            val_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    # Calculate metrics
    train_loss = running_loss / len(train_dataset)
    val_loss = val_loss / len(val_dataset)
    val_acc = 100 * correct / total
    
    # Update scheduler
    scheduler.step(val_acc)
    
    # Save best model
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best_inceptionv3.pth')
    
    print(f'Epoch {epoch+1}/50 | '
          f'Train Loss: {train_loss:.4f} | '
          f'Val Loss: {val_loss:.4f} | '
          f'Val Acc: {val_acc:.2f}% | '
          f'LR: {optimizer.param_groups[0]["lr"]:.2e}')