Usage

Training Models

Once dependencies are installed, you can start training the models.

Navigate to the galaxy directory:

cd galaxyHackers/galaxy

Train a model by running the main.py script:

python3 main.py --models MODEL_NAME --epochs NUM_EPOCH --data DATASET

Available flags:

  • --models: The name of the model(s) you wish to train.
    Available options:
    • Baseline
    • ResNet18
    • EfficientNet
    • DenseNet
    • SpinalNet_ResNet
    • SpinalNet_VGG
    • ViTL16
    • AlexNet_VGG
    • CNN_MLP
    • YOLO
    Default: Baseline, ResNet18, EfficientNet, DenseNet, SpinalNet_ResNet, SpinalNet_VGG, ViTL16, AlexNet_VGG, CNN_MLP, YOLO
  • --epochs: The number of epochs for training.
    Default: 5
  • --mm: The momentum you wish to use in your optimizer (relevant if you use SGD or RMSprop).
    Default: 0.9
  • --optimizer: The optimizer you wish to use in training.
    Available options:
    • Adam
    • SGD
    • Rprop
    • NAdam
    • RAdam
    • AdamW
    • RMSprop
    • DiffGrad
    Default: Adam
  • --segment: Pass this flag, if you want to apply the trained models to get segmentation maps.
  • --data: The dataset that you want to download for training. In case script fails to download dataset, it will address unWISE survey website (available with VPN only).
    Available options:
    • WISE
    • ACT
    Default: WISE

Learning rate is not passed as a parameter because the optimal learning rate is found automatically.

Running Multiple Models Simultaneously

You can run several models simultaneously with the same optimizer and learning rate.

Note

The script does not support using different optimizers for each model at this time.

Example Commands

Single model training:

python3 main.py --models YOLO --epochs 5 --data ACT --segment

Multiple models training:

python3 main.py --models AlexNet_VGG ResNet18 --epochs 20 --data WISE --segment
python3 main.py --models --epochs 20 --data WISE --segment

Output and Logs

  • Metrics: Training and validation metrics are logged to Comet ML and saved as CSV files (*_train_metrics.csv, *_val_metrics.csv).

  • Segmentation Plots: Generated for each model and stored locally.

  • Model Performance: Combines metrics across models for analysis.