train.py

import torch
import torch.nn as nn
import torch.optim as optim
from model import CustomCNN
from data import get_data_loaders
from tqdm import tqdm
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime

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

# Veri yükle
train_loader, val_loader, class_names = get_data_loaders('Flowers/train', batch_size=32)

# Modeli oluştur
model = CustomCNN(num_classes=len(class_names)).to(device)

def train_model(model, train_loader, val_loader, num_epochs=10, learning_rate=0.001, device='cuda'):
    """Train the model and return training history"""
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=2)
    
    history = {
        'train_loss': [],
        'val_loss': [],
        'train_acc': [],
        'val_acc': [],
        'train_precision': [],
        'val_precision': [],
        'train_recall': [],
        'val_recall': [],
        'train_f1': [],
        'val_f1': []
    }
    
    best_val_acc = 0.0
    start_time = datetime.now()
    
    for epoch in range(num_epochs):
        # Training phase
        model.train()
        train_loss = 0.0
        train_preds = []
        train_labels = []
        
        for inputs, labels in tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs} [Train]'):
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, preds = torch.max(outputs, 1)
            train_preds.extend(preds.cpu().numpy())
            train_labels.extend(labels.cpu().numpy())
        
        # Calculate training metrics
        train_loss = train_loss / len(train_loader)
        train_acc = accuracy_score(train_labels, train_preds)
        train_precision, train_recall, train_f1, _ = precision_recall_fscore_support(
            train_labels, train_preds, average='weighted'
        )
        
        # Validation phase
        model.eval()
        val_loss = 0.0
        val_preds = []
        val_labels = []
        
        with torch.no_grad():
            for inputs, labels in tqdm(val_loader, desc=f'Epoch {epoch+1}/{num_epochs} [Val]'):
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item()
                _, preds = torch.max(outputs, 1)
                val_preds.extend(preds.cpu().numpy())
                val_labels.extend(labels.cpu().numpy())
        
        # Calculate validation metrics
        val_loss = val_loss / len(val_loader)
        val_acc = accuracy_score(val_labels, val_preds)
        val_precision, val_recall, val_f1, _ = precision_recall_fscore_support(
            val_labels, val_preds, average='weighted'
        )
        
        # Update learning rate
        scheduler.step(val_acc)
        
        # Save best model
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), f'best_model_{model.__class__.__name__}.pth')
        
        # Update history
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)
        history['train_precision'].append(train_precision)
        history['val_precision'].append(val_precision)
        history['train_recall'].append(train_recall)
        history['val_recall'].append(val_recall)
        history['train_f1'].append(train_f1)
        history['val_f1'].append(val_f1)
        
        print(f'Epoch {epoch+1}/{num_epochs}:')
        print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}')
        print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')
    
    end_time = datetime.now()
    training_time = (end_time - start_time).total_seconds()
    
    return history, training_time

def visualize_features(model, data_loader, device='cuda', save_path=None):
    """Visualize feature maps from different layers"""
    model.eval()
    inputs, _ = next(iter(data_loader))
    inputs = inputs.to(device)
    
    # Forward pass to get feature maps
    _ = model(inputs)
    
    # Get feature maps
    feature_maps = model.feature_maps
    
    # Create visualization
    fig, axes = plt.subplots(len(feature_maps), 1, figsize=(10, 5*len(feature_maps)))
    if len(feature_maps) == 1:
        axes = [axes]
    
    for ax, (layer_name, features) in zip(axes, feature_maps.items()):
        # Select first 16 feature maps
        features = features[0, :16].cpu()
        
        # Create grid of feature maps
        grid = torch.zeros((4, 4, features.shape[1], features.shape[2]))
        for i in range(16):
            row = i // 4
            col = i % 4
            grid[row, col] = features[i]
        
        # Normalize for visualization
        grid = (grid - grid.min()) / (grid.max() - grid.min())
        
        # Plot
        ax.imshow(grid.permute(0, 2, 1, 3).reshape(4*features.shape[1], 4*features.shape[2]),
                 cmap='viridis')
        ax.set_title(f'Feature Maps from {layer_name}')
        ax.axis('off')
    
    plt.tight_layout()
    if save_path:
        plt.savefig(save_path)
    plt.close()

def plot_training_history(history, save_path=None):
    """Plot training history"""
    metrics = ['loss', 'acc', 'precision', 'recall', 'f1']
    fig, axes = plt.subplots(len(metrics), 1, figsize=(10, 5*len(metrics)))
    
    for ax, metric in zip(axes, metrics):
        train_metric = history[f'train_{metric}']
        val_metric = history[f'val_{metric}']
        
        ax.plot(train_metric, label=f'Train {metric.capitalize()}')
        ax.plot(val_metric, label=f'Val {metric.capitalize()}')
        ax.set_xlabel('Epoch')
        ax.set_ylabel(metric.capitalize())
        ax.legend()
        ax.grid(True)
    
    plt.tight_layout()
    if save_path:
        plt.savefig(save_path)
    plt.close()

def evaluate_model(model, test_loader, device='cuda'):
    """Evaluate model on test set and return metrics"""
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in tqdm(test_loader, desc='Evaluating'):
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # Calculate metrics
    accuracy = accuracy_score(all_labels, all_preds)
    precision, recall, f1, _ = precision_recall_fscore_support(
        all_labels, all_preds, average='weighted'
    )
    
    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1
    }