main.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms

from models.convnext import ConvNeXt
from models.resnet import ResNet
from models.vit import VisionTransformer
from utils.data import FishDataset, create_data_loaders, get_num_classes_dict
from utils.train import compute_hierarchical_accuracy, save_model, train_hierarchical_model
from utils.plot import plot_training_curves
from models.cnn import CNN
from models.dnn import DNN

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

def main():
    print("="*60)
    print("HIERARCHICAL FISH CLASSIFICATION TRAINING")
    print("="*60)
    
    print("\nLoading datasets...")
    # hyperparameters
    batch_size = 64
    num_workers = 4
    train_dataset_fraction = 1.0
    scale_loss_weights = True
    # taxonomic_levels = ['Class', 'Order', 'Family', 'Genus', 'species']
    taxonomic_levels = ['species']
    
    print(f"\nTraining dataset fraction: {train_dataset_fraction}")
    train_dataset = FishDataset(
        csv_file='data/processed/train_split.csv',
        root_dir='data/Image_Library',
        transform=transform,
        taxonomic_levels=taxonomic_levels
    )
    
    val_dataset = FishDataset(
        csv_file='data/processed/val_split.csv',
        root_dir='data/Image_Library',
        transform=transform,
        taxonomic_levels=taxonomic_levels
    )
    
    test_dataset = FishDataset(
        csv_file='data/processed/test_split.csv',
        root_dir='data/Image_Library',
        transform=transform,
        taxonomic_levels=taxonomic_levels
    )
    
    print(f"Train set: {len(train_dataset)} samples")
    print(f"Validation set: {len(val_dataset)} samples")
    print(f"Test set: {len(test_dataset)} samples")
    
    print("\nNumber of classes per taxonomic level:")
    num_classes_dict = get_num_classes_dict([train_dataset, val_dataset, test_dataset], taxonomic_levels)
    for level in taxonomic_levels:
        print(f"  {level}: {num_classes_dict[level.lower()]}")
    
    
    # inversely scaled loss weights (more labels -> higher weight)
    max_num_classes = max(num_classes_dict.values()) if num_classes_dict.values() else 1
    loss_weights = {}
    for level, num_classes in num_classes_dict.items():
        if num_classes > 0 and scale_loss_weights:
            loss_weights[level] = num_classes / max_num_classes
        else:
            loss_weights[level] = 1.0
    
    print("\nLoss weights (scaled by number of classes):")
    for level, weight in loss_weights.items():
        print(f"  {level}: {weight:.4f} (num_classes: {num_classes_dict[level]})")
    
    # create data loaders
    loaders = create_data_loaders(
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        test_dataset=test_dataset,
        train_dataset_fraction=train_dataset_fraction,
        batch_size=batch_size,
        num_workers=num_workers
    )
    train_loader = loaders['train']
    val_loader = loaders['val']
    

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nUsing device: {device}")
    
    def run_experiment(model, num_epochs, learning_rate, model_name):
        print("\n" + "="*60)
        print(f"TRAINING {model_name} MODEL")
        print("="*60)
        
        model = model.to(device)
        optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
        criterion = nn.CrossEntropyLoss()
        
        (train_losses_total, train_losses_per_level,
        train_accuracies_total, train_accuracies_per_level,
        eval_accuracies_total, eval_accuracies_per_level) = train_hierarchical_model(
            net=model,
            train_loader=train_loader,
            eval_loader=val_loader, 
            num_epochs=num_epochs,
            optimizer=optimizer,
            criterion=criterion,
            loss_weights=loss_weights,
            max_eval_batches=100
        )
        
        test_acc_total, test_acc_per_level = compute_hierarchical_accuracy(model, val_loader, list(num_classes_dict.keys()))
        print(f"Test accuracy: {test_acc_total:.4f}")
        print(f"Test accuracy per level: {test_acc_per_level}")
        
        save_model(model, f'model_files/{model_name.lower()}.pth')
        
        plot_training_curves(
            model_name,
            train_losses_total, train_losses_per_level,
            train_accuracies_total, train_accuracies_per_level,
            eval_accuracies_total, eval_accuracies_per_level,
            list(num_classes_dict.keys()),
            final_test_accuracy_total=test_acc_total,
            final_test_accuracies_per_level=test_acc_per_level
        )
        
        print("\n" + "="*60)
        print("TRAINING COMPLETE!")
        print("="*60)

        return test_acc_total, test_acc_per_level
    
    cnn_model = CNN(num_classes_dict=num_classes_dict)
    run_experiment(model = cnn_model, num_epochs = 100, learning_rate = 0.0001, model_name = 'CNN-genus')
    
    dnn_model = DNN(num_classes_dict=num_classes_dict)
    run_experiment(model = dnn_model, num_epochs = 100, learning_rate = 0.0001, model_name = 'DNN-genus')
    
    vit_model = VisionTransformer(num_classes_dict=num_classes_dict)
    run_experiment(model = vit_model, num_epochs = 100, learning_rate = 0.00001, model_name = 'Pre-Trained-VIT-species')
    
    resnet_model = ResNet(num_classes_dict=num_classes_dict)
    run_experiment(model = resnet_model, num_epochs = 100, learning_rate = 0.0001, model_name = 'Pre-Trained-ResNet50-genus')
    
    convnext_model = ConvNeXt(num_classes_dict=num_classes_dict)
    run_experiment(model = convnext_model, num_epochs = 100, learning_rate = 0.00005, model_name = 'Pre-Trained-ConvNext-genus')
    

if __name__ == '__main__':
    main()