Hierarchical Multi-Task Learning: Leveraging Taxonomic Structure for Fish Classification

Research Overview

This repository implements hierarchical multi-task learning models for fish classification using the FishNet dataset. We investigate multi-output neural networks that simultaneously predict multiple taxonomic levels (class, order, family, genus, and species) and compare them against ensembles of single-output models. Our best hierarchical model (ConvNeXt-Tiny with equal loss weighting) achieves 98.63% class, 88.83% order, 83.81% family, 75.54% genus, and 25.26% species accuracy, representing improvements of 20.58% and 13.06% over the FishNet benchmark at family and order levels respectively.

Key Finding: Hierarchical multi-output models with shared representations outperform single-output ensemble approaches in both per-level accuracy and total accuracy metrics. The best multi-output model achieves 74.42% total accuracy compared to 71.54% for the ensemble approach using the same model architecture and training routine.

For detailed methodology, results, and analysis, see the paper.

Dataset: This project uses the FishNet dataset containing 94,532 images across 17,357 species. Khan et al. established benchmarks for fish classification in their paper "FishNet: A Large-scale Dataset and Benchmark for Fish Recognition, Detection, and Functional Traits Prediction." Both the FishNet dataset and benchmark are available there.

Project Structure

fish-classification/
├── models/          # Model architectures (DNN, CNN, ResNet, ConvNeXt, ViT)
├── utils/           # Data loading, training, evaluation, and plotting utilities
├── data/            # Dataset files and processed splits
├── graphs/          # Training curves and visualizations
├── model_files/     # Saved model checkpoints
├── test.py          # Evaluation script
├── run_all_tests.py # Script to run tests for all models and configurations
└── main.py          # Main training script

Models

This repository implements 5 model architectures, all supporting hierarchical multi-output classification:

• DNN: Simple baseline with fully connected layers
• CNN: Custom convolutional neural network
• ResNet-50: Pre-trained ResNet with transfer learning
• ConvNeXt-Tiny: Pre-trained ConvNeXt
• ViT-Base: Pre-trained Vision Transformer

All models use a shared backbone architecture with separate classification heads for each taxonomic level. Two loss weighting strategies are supported: equal weighting (treats all levels uniformly) and scaled weighting (adjusts contribution proportionally to number of classes).

Data Preprocessing

preprocess_data.py performs the following actions:

• Combines pre-made data/raw/train.csv and data/raw/test.csv into a single dataset
• Removes rows that are missing a label in the stratification column (Genus by default)
  • (Genus is the default stratification column to ensure enough samples exist to perform stratification while using the lowest possible taxonomic class)
• Filters out rare Genus classes (required for effective stratification)
• Creates integer label encodings for all taxonomic levels:
  • Class
  • Order
  • Family
  • Genus
  • Species
• Performs stratified train/val/test split (70%/15%/15%) based on Genus
• Saves three new CSV files:
  • data/processed/train_split.csv (70% of data)
  • data/processed/val_split.csv (15% of data)
  • data/processed/test_split.csv (15% of data)
• Saves label mappings to data/processed/label_mappings.json
  • Missing labels are given a -1 label mapping (the dataset has missing species values for about 60% of samples)

Quickstart

Prerequisites

Download the fish image dataset as well as the labels into the /data folder

The /data directory should look like this:

data/
├── Image_Library/
│   └── ...
├── raw/
│   ├── train.csv
│   ├── test.csv
│   └── ...

1. Run the preprocessing script using preprocess-data.py

After running preprocessing, you'll have:

data/
├── raw/
│   ├── train.csv
│   ├── test.csv
│   └── ...
├── processed/
│   ├── train_split.csv
│   ├── val_split.csv
│   ├── test_split.csv
│   └── label_mappings.json
└── Image_Library/

2. Train your models using main.py:

   # Edit main.py to configure your models and training parameters
   python main.py

See main.py for examples of training hierarchical multi-output models and single-output models.

3. Evaluate your models using test.py:

   You can evaluate a trained model using test.py with command line arguments. For a multi-output (hierarchical) model, run, for example:

   python test.py --mode single --model_path model_files/hierachical_output/flat_loss/vit.pth --model_class vit

   For an ensemble of five single-output models (one for each taxonomic level), run:

   python test.py --mode ensemble \
      --ensemble_paths model_files/single_output/resnet50/resnet50-class.pth \
                     model_files/single_output/resnet50/resnet50-order.pth \
                     model_files/single_output/resnet50/resnet50-family.pth \
                     model_files/single_output/resnet50/resnet50-genus.pth \
                     model_files/single_output/resnet50/resnet50-species.pth \
      --ensemble_model_class resnet

   See the test.py file for all command line arguments and further instructions.

   Alternatively, you can run all model evaluations automatically using run_all_tests.py. This script will search for compatible model checkpoint files within the model_files/hierachical_output/ (for multi-output models) and model_files/single_output/ (for single-output/ensemble models) directory structures, and run all possible evaluations. Make sure your model checkpoints are organized according to this standard directory structure:

   model_files/
   ├── hierachical_output/
   │   └── <loss_type>/
   │       ├── vit.pth
   │       ├── resnet.pth
   │       └── ...
   └── single_output/
       └── <model_type>/
           ├── <model_type>-class.pth
           ├── <model_type>-order.pth
           ├── <model_type>-family.pth
           ├── <model_type>-genus.pth
           └── <model_type>-species.pth
   

   Where loss_type can be either flat_loss or scaled_loss and model_type can be any of the following: cnn, dnn, resnet50, vit, or convnext

   Running python run_all_tests.py will automatically evaluate all models in these folders.