Fix errors and bugs in the given code, without changing its intended functionality. ONLY return the fixed code and nothing else. Enclose the entire code in a single code block.



import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms, datasets, models
from torchvision.models import mobilenet_v3_small, squeezenet1_1, alexnet, regnet_x_800mf
from sklearn.metrics import classification_report, accuracy_score, precision_recall_fscore_support, confusion_matrix
import numpy as np
import os
from tqdm import tqdm
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Set random seed for reproducibility
torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create directories
checkpoint_dir = 'checkpoints'
results_dir = 'results'
plots_dir = 'plots'
os.makedirs(checkpoint_dir, exist_ok=True)
os.makedirs(results_dir, exist_ok=True)
os.makedirs(plots_dir, exist_ok=True)

# Data transforms
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def plot_confusion_matrix(y_true, y_pred, classes, model_name):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=classes, yticklabels=classes)
    plt.title(f'Confusion Matrix - {model_name}')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, f'{model_name}_confusion_matrix.png'))
    plt.close()

def plot_learning_curves(history, model_name):
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 2, 1)
    plt.plot(history['train_loss'], label='Train Loss')
    plt.plot(history['val_loss'], label='Val Loss')
    plt.title(f'{model_name} - Loss Curves')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    
    plt.subplot(1, 2, 2)
    plt.plot(history['val_acc'], label='Val Accuracy')
    plt.title(f'{model_name} - Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, f'{model_name}_learning_curves.png'))
    plt.close()

def load_data(data_dir, batch_size=32):
    train_dir = os.path.join(data_dir, 'train')
    val_dir = os.path.join(data_dir, 'val')
    
    train_dataset = datasets.ImageFolder(train_dir, transform=transform)
    val_dataset = datasets.ImageFolder(val_dir, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    
    return train_loader, val_loader, train_dataset.classes

def initialize_model(model_name, num_classes):
    if model_name == 'mobilenet_v3s':
        model = mobilenet_v3_small(pretrained=True)
        model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, num_classes)
    elif model_name == 'squeezenet':
        model = squeezenet1_1(pretrained=True)
        model.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=1)
    elif model_name == 'alexnet':
        model = alexnet(pretrained=True)
        model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_classes)
    add regnet_x_800mf...
    
    return model.to(device)

def load_checkpoint(model, optimizer, model_name):
    checkpoint_path = os.path.join(checkpoint_dir, f'{model_name}_checkpoint.pth')
    if os.path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint['epoch']
        best_val_acc = checkpoint['val_acc']
        print(f"Resuming training from epoch {start_epoch}")
        return start_epoch, best_val_acc
    return 0, 0.0

def save_checkpoint(model, optimizer, epoch, val_acc, model_name):
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'val_acc': val_acc,
    }
    checkpoint_path = os.path.join(checkpoint_dir, f'{model_name}_checkpoint.pth')
    torch.save(checkpoint, checkpoint_path)

def quick_evaluate(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    running_loss = 0
    criterion = nn.CrossEntropyLoss()
    
    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    accuracy = 100. * correct / total
    avg_loss = running_loss / len(data_loader)
    return accuracy, avg_loss

def final_evaluate(model, data_loader, classes, model_name):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in tqdm(data_loader, desc='Final Evaluation'):
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # Calculate metrics
    accuracy = accuracy_score(all_labels, all_preds)
    precision, recall, f1, support = precision_recall_fscore_support(all_labels, all_preds, average='weighted')
    report = classification_report(all_labels, all_preds, target_names=classes, digits=4)
    
    # Plot confusion matrix
    plot_confusion_matrix(all_labels, all_preds, classes, model_name)
    
    metrics = {
        'overall': {
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1': f1
        },
        'raw': {
            'true_labels': all_labels,
            'predictions': all_preds
        },
        'report': report
    }
    
    return metrics

def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, model_name):
    history = {
        'train_loss': [],
        'val_loss': [],
        'val_acc': []
    }
    
    start_epoch, best_val_acc = load_checkpoint(model, optimizer, model_name)
    
    for epoch in range(start_epoch, num_epochs):
        # Training phase
        model.train()
        running_loss = 0.0
        
        for inputs, labels in tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}'):
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
        
        # Quick validation
        val_acc, val_loss = quick_evaluate(model, val_loader)
        avg_train_loss = running_loss / len(train_loader)
        
        # Save metrics
        history['train_loss'].append(avg_train_loss)
        history['val_loss'].append(val_loss)
        history['val_acc'].append(val_acc)
        
        print(f'Epoch {epoch+1}/{num_epochs}:')
        print(f'Train Loss: {avg_train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
        
        # Save checkpoint if better validation accuracy
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            save_checkpoint(model, optimizer, epoch + 1, val_acc, model_name)
            torch.save(model.state_dict(), os.path.join(checkpoint_dir, f'best_{model_name}_model.pth'))
            print(f'New best accuracy! Checkpoint saved.')
    
    # Save learning curves
    df = pd.DataFrame({
        'epoch': range(1, num_epochs + 1),
        'train_loss': history['train_loss'],
        'val_loss': history['val_loss'],
        'val_accuracy': history['val_acc']
    })
    df.to_csv(os.path.join(results_dir, f'{model_name}_learning_curves.csv'), index=False)
    
    # Plot learning curves
    plot_learning_curves(history, model_name)
    
    return history

def main():
    data_dir = '/kaggle/working/data'
    batch_size = 32
    num_epochs = 100
    learning_rate = 0.001
    
    train_loader, val_loader, classes = load_data(data_dir, batch_size)
    num_classes = len(classes)
    
    models_to_train = {
        # 'mobilenet_v3s': mobilenet_v3_small,
        # 'squeezenet': squeezenet1_1,
        'alexnet': alexnet,
        'regnet' regnet_x_800mf
    }
    
    # Create results summary dataframe
    results_summary = []
    
    for model_name in models_to_train:
        print(f'\n{"="*50}')
        print(f'Training {model_name}...')
        print(f'{"="*50}')
        
        model = initialize_model(model_name, num_classes)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        
        # Train model
        history = train_model(model, train_loader, val_loader, criterion, 
                          optimizer, num_epochs, model_name)
        
        # Load best model for final evaluation
        best_model_path = os.path.join(checkpoint_dir, f'best_{model_name}_model.pth')
        model.load_state_dict(torch.load(best_model_path))
        
        # Final evaluation
        print(f'\nPerforming final evaluation for {model_name}...')
        metrics = final_evaluate(model, val_loader, classes, model_name)
        
        # Save results to summary
        results_summary.append({
            'model': model_name,
            'accuracy': metrics['overall']['accuracy'],
            'precision': metrics['overall']['precision'],
            'recall': metrics['overall']['recall'],
            'f1': metrics['overall']['f1']
        })
        
        # Print metrics
        print(f'\nFinal Metrics for {model_name}:')
        print(f'Accuracy: {metrics["overall"]["accuracy"]:.4f}')
        print(f'Precision: {metrics["overall"]["precision"]:.4f}')
        print(f'Recall: {metrics["overall"]["recall"]:.4f}')
        print(f'F1-score: {metrics["overall"]["f1"]:.4f}')
        print('\nClassification Report:')
        print(metrics['report'])
        print(f'{"="*50}\n')
    
    # Save overall results summary
    pd.DataFrame(results_summary).to_csv(os.path.join(results_dir, 'model_comparison.csv'), index=False)

if __name__ == '__main__':
    main()