train.py

import os
import torch
import numpy
from datetime import datetime
from PIL import ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

# Import original modules with aliases
import util as toolkit
from loader import get_loader as fetch_train_data, get_val_loader as fetch_val_data
from config import ConfigurationManager as Configurator
from model import model as NeuralNetwork
from util import bceLoss as compute_binary_loss


def prepare_validation_config():
    """Create validation-specific configuration"""
    val_cfg = Configurator().parse()
    val_cfg.isTrain = False
    val_cfg.isVal = True

    return val_cfg


def execute_training_iteration(
        data_provider,
        network,
        optimizer,
        epoch_index,
        storage_location
):
    """Perform training iteration"""
    network.train()
    global iteration_counter
    epoch_iterations = 0
    total_loss = 0

    try:
        for batch_idx, (inputs, targets) in enumerate(data_provider, start=1):
            optimizer.zero_grad()

            # Move data to GPU
            inputs = inputs.cuda()
            targets = targets.cuda()

            # Forward pass
            outputs = network(inputs).ravel()

            # Compute loss
            loss_function = compute_binary_loss()
            batch_loss = loss_function(outputs, targets)

            # Backward pass
            batch_loss.backward()
            optimizer.step()

            # Update counters
            iteration_counter += 1
            epoch_iterations += 1
            total_loss += batch_loss.item()

            # Log progress
            if batch_idx % 500 == 0 or batch_idx == total_batches or batch_idx == 1:
                current_time = datetime.now().strftime("%Y%m%d-%H%M%S")
                progress_percent = (batch_idx / total_batches) * 100

                status_report = (
                    f"📊 Epoch: {epoch_index:02d}/{config.epoch:02d} | "
                    f"🔢 Iteration: {batch_idx:04d}/{total_batches:04d} "
                    f"({progress_percent:.1f}%) | "
                    f"📉 Loss Metric: {batch_loss.item():.6f}"
                )
                print(status_report)

        # Save periodic checkpoint
        if epoch_index % 50 == 0:
            checkpoint_path = os.path.join(
                storage_location,
                f'Network_epoch_{epoch_index}.pth'
            )
            torch.save(network.state_dict(), checkpoint_path)

    except KeyboardInterrupt:
        print("Training interrupted: saving model and exiting")


def perform_validation(
        validation_sets,
        network,
        epoch_index,
        storage_location
):
    """Evaluate model on validation sets"""
    network.eval()
    global best_performing_epoch, highest_accuracy

    total_correct = total_samples = 0

    with torch.no_grad():
        for dataset in validation_sets:
            correct_ai = correct_nature = 0

            name = dataset['name']
            ai_loader = dataset['val_ai_loader']
            ai_count = dataset['ai_size']
            nature_loader = dataset['val_nature_loader']
            nature_count = dataset['nature_size']

            print(f"||Validating||")

            # Process AI-generated images
            for inputs, targets in ai_loader:
                inputs = inputs.cuda()
                targets = targets.cuda()

                predictions = network(inputs)
                probabilities = torch.sigmoid(predictions).ravel()

                # Count correct predictions
                correct = (
                        ((probabilities > 0.5) & (targets == 1)) |
                        ((probabilities < 0.5) & (targets == 0))
                )
                correct_ai += correct.sum().item()

            ai_accuracy = correct_ai / ai_count
            #print(f"AI Accuracy: {ai_accuracy:.4f}")

            # Process natural images
            for inputs, targets in nature_loader:
                inputs = inputs.cuda()
                targets = targets.cuda()

                predictions = network(inputs)
                probabilities = torch.sigmoid(predictions).ravel()

                correct = (
                        ((probabilities > 0.5) & (targets == 1)) |
                        ((probabilities < 0.5) & (targets == 0))
                )
                correct_nature += correct.sum().item()

            nature_accuracy = correct_nature / nature_count
            #print(f"Nature Accuracy: {nature_accuracy:.4f}")

            # Calculate dataset accuracy
            dataset_accuracy = (correct_ai + correct_nature) / (ai_count + nature_count)
            total_correct += correct_ai + correct_nature
            total_samples += ai_count + nature_count

            print(f"Epoch: {epoch_index}, Accuracy: {dataset_accuracy:.4f}")

    # Calculate overall accuracy
    overall_accuracy = total_correct / total_samples

    # Save best model
    if epoch_index == 1:
        best_performing_epoch = 1
        highest_accuracy = overall_accuracy
        best_model_path = os.path.join(storage_location, 'Network_best.pth')
        torch.save(network.state_dict(), best_model_path)
        print(f"Saved best model on Epoch: {epoch_index}")
    else:
        if overall_accuracy > highest_accuracy:
            best_performing_epoch = epoch_index
            highest_accuracy = overall_accuracy
            best_model_path = os.path.join(storage_location, 'Network_best.pth')
            torch.save(network.state_dict(), best_model_path)
            print(f"Saved best model on Epoch: {epoch_index}")

    print(
        f"🏆 Performance Report | "
        f"Current Epoch: {epoch_index:03d} | "
        f"Accuracy Score: {overall_accuracy:.2%} | "
        f"Peak Performance: Epoch {best_performing_epoch:03d} | "
        f"Highest Accuracy: {highest_accuracy:.2%}"
    )


def configure_gpu(gpu_id):
    """Set GPU configuration"""
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id


def main_execution():
    """Main training procedure"""
    # Initialize environment
    torch.set_num_threads(2)
    toolkit.set_random_seed()

    # Load configurations
    global config
    config = Configurator().parse()
    val_config = prepare_validation_config()

    # Prepare data
    global total_batches
    train_loader = fetch_train_data(config)
    total_batches = len(train_loader)
    val_loader = fetch_val_data(val_config)

    # Configure GPU
    configure_gpu(config.gpu_id)

    # Initialize model
    model = NeuralNetwork().cuda()
    if config.load:
        model.load_state_dict(torch.load(config.load))
        print(f"Loaded model from {config.load}")

    # Prepare optimizer
    optimizer = torch.optim.Adam(model.parameters(), config.lr)

    # Create output directory
    output_dir = config.save_path
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Initialize training state
    global iteration_counter, best_performing_epoch, highest_accuracy
    iteration_counter = 0
    best_performing_epoch = 0
    highest_accuracy = 0

    print("||Training||")

    # Training loop
    for epoch in range(1, config.epoch + 1):
        # Adjust learning rate
        current_lr = toolkit.poly_lr(optimizer, config.lr, epoch, config.epoch)

        # Training iteration
        execute_training_iteration(
            train_loader, model, optimizer, epoch, output_dir
        )

        # Validation
        perform_validation(
            val_loader, model, epoch, output_dir
        )


if __name__ == '__main__':
    main_execution()