zoobot 1.0.4

Galaxy morphology classifiers

Zoobot

Zoobot classifies galaxy morphology with deep learning.

Zoobot is trained using millions of answers by Galaxy Zoo volunteers. This code will let you retrain Zoobot to accurately solve your own prediction task.

• Install
• Quickstart
• Worked Examples
• Pretrained Weights
• Datasets
• Documentation (for understanding/reference)

Installation

You can retrain Zoobot in the cloud with a free GPU using this Google Colab notebook. To install locally, keep reading.

Download the code using git:

git clone [e-mail]:mwalmsley/zoobot.git

And then pick one of the three commands below to install Zoobot and either PyTorch (recommended) or TensorFlow:

# Zoobot with PyTorch and a GPU. Requires CUDA 11.3.
pip install -e "zoobot[pytorch_cu113]" --extra-index-url https://download.pytorch.org/whl/cu113

# OR Zoobot with PyTorch and no GPU
pip install -e "zoobot[pytorch_cpu]" --extra-index-url https://download.pytorch.org/whl/cpu

# OR Zoobot with PyTorch on Mac with M1 chip
pip install -e "zoobot[pytorch_m1]"

# OR Zoobot with TensorFlow. Works with and without a GPU, but if you have a GPU, you need CUDA 11.2. 
pip install -e "zoobot[tensorflow]

This installs the downloaded Zoobot code using pip editable mode so you can easily change the code locally. Zoobot is also available directly from pip (pip install zoobot[option]). Only use this if you are sure you won't be making changes to Zoobot itself. For Google Colab, use pip install zoobot[pytorch_colab]

To use a GPU, you must already have CUDA installed and matching the versions above. I share my install steps here. GPUs are optional - Zoobot will run retrain fine on CPU, just slower.

Quickstart

The Colab notebook is the quickest way to get started. Alternatively, the minimal example below illustrates how Zoobot works.

Let's say you want to find ringed galaxies and you have a small labelled dataset of 500 ringed or not-ringed galaxies. You can retrain Zoobot to find rings like so:

    import pandas as pd
    from galaxy_datasets.pytorch.galaxy_datamodule import GalaxyDataModule
    from zoobot.pytorch.training import finetune

    # csv with 'ring' column (0 or 1) and 'file_loc' column (path to image)
    labelled_df = pd.read_csv('/your/path/some_labelled_galaxies.csv')

    datamodule = GalaxyDataModule(
      label_cols=['ring'],
      catalog=labelled_df,
      batch_size=32
    )

    # load trained Zoobot model
    model = finetune.FinetuneableZoobotClassifier(checkpoint_loc, num_classes=2)  
    
    # retrain to find rings
    trainer = finetune.get_trainer(save_dir)
    trainer.fit(model, datamodule)

Then you can make predict if new galaxies have rings:

    from zoobot.pytorch.predictions import predict_on_catalog

    # csv with 'file_loc' column (path to image). Zoobot will predict the labels.
    unlabelled_df = pd.read_csv('/your/path/some_unlabelled_galaxies.csv')

    predict_on_catalog.predict(
      unlabelled_df,
      model,
      label_cols=['ring'],  # only used for 
      save_loc='/your/path/finetuned_predictions.csv'
    )

Zoobot includes many guides and working examples - see the Getting Started section below.

Getting Started

I suggest starting with the Colab notebook or the worked examples below, which you can copy and adapt.

For context and explanation, see the documentation.

For pretrained model weights, precalculated representations, catalogues, and so forth, see the data notes in particular.

Worked Examples

PyTorch (recommended):

• pytorch/examples/finetuning/finetune_binary_classification.py
• pytorch/examples/finetuning/finetune_counts_full_tree.py
• pytorch/examples/representations/get_representations.py
• pytorch/examples/train_model_on_catalog.py (only necessary to train from scratch)

TensorFlow:

• tensorflow/examples/train_model_on_catalog.py (only necessary to train from scratch)
• tensorflow/examples/make_predictions.py
• tensorflow/examples/finetune_minimal.py
• tensorflow/examples/finetune_advanced.py

There is more explanation and an API reference on the docs.

I also include the scripts used to create and benchmark our pretrained models. Many pretrained models are available already, but if you need one trained on e.g. different input image sizes or with a specific architecture, I can probably make it for you.

When trained with a decision tree head (ZoobotTree, FinetuneableZoobotTree), Zoobot can learn from volunteer labels of varying confidence and predict posteriors for what the typical volunteer might say. Specifically, this Zoobot mode predicts the parameters for distributions, not simple class labels! For a demonstration of how to interpret these predictions, see the gz_decals_data_release_analysis_demo.ipynb.

(Optional) Install PyTorch or TensorFlow, with CUDA

If you're not using a GPU, skip this step. Use the pytorch_cpu or tensorflow_cpu options in the section below.

Install PyTorch 1.12.1 or Tensorflow 2.10.0 and compatible CUDA drivers. I highly recommend using conda to do this. Conda will handle both creating a new virtual environment (conda create) and installing CUDA (cudatoolkit, cudnn)

CUDA 11.3 for PyTorch:

conda create --name zoobot38_torch python==3.8
conda activate zoobot38_torch
conda install -c conda-forge cudatoolkit=11.3

CUDA 11.2 and CUDNN 8.1 for TensorFlow 2.10.0:

conda create --name zoobot38_tf python==3.8
conda activate zoobot38_tf
conda install -c conda-forge cudatoolkit=11.2 cudnn=8.1.0
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/  # add this environment variable

Latest minor features (v1.0.4)

• Now supports multi-class finetuning. See pytorch/examples/finetuning/finetune_multiclass_classification.py
• Removed simplejpeg dependency due to M1 install issue.
• Pinned timm version to ensure MaX-ViT models load correctly. Models supporting the latest timm will follow.
• (internal until published) GZ Evo v2 now includes Cosmic Dawn (HSC). Significant performance improvement on HSC finetuning.

Latest major features (v1.0.0)

v1.0.0 recognises that most of the complexity in this repo is training Zoobot from scratch, but most non-GZ users will probably simply want to load the pretrained Zoobot and finetune it on their data.

• Adds new finetuning interface (finetune.run_finetuning()), examples.
• Refocuses docs on finetuning rather than training from scratch.
• Rework installation process to separate CUDA from Zoobot (simpler, easier)
• Better wandb logging throughout, to monitor training
• Remove need to make TFRecords. Now TF directly uses images.
• Refactor out augmentations and datasets to galaxy-datasets repo. TF and Torch now use identical augmentations (via albumentations).
• Many small quality-of-life improvements

Contributions are very welcome and will be credited in any future work. Please get in touch! See CONTRIBUTING.md for more.

Benchmarks and Replication - Training from Scratch

The benchmarks folder contains slurm and Python scripts to train Zoobot from scratch. We use these scripts to make sure new code versions work well, and that TensorFlow and PyTorch achieve similar performance.

Training Zoobot using the GZ DECaLS dataset option will create models very similar to those used for the GZ DECaLS catalogue and shared with the early versions of this repo. The GZ DESI Zoobot model is trained on additional data (GZD-1, GZD-2), as the GZ Evo Zoobot model (GZD-1/2/5, Hubble, Candels, GZ2).

Citing

We have submitted a JOSS paper to describe Zoobot itself. We hope this will become the single point-of-reference for Zoobot. Meanwhile, please cite the Galaxy Zoo DECaLS, which uses the code that evolved into Zoobot:

@article{Walmsley2022decals,
author = {Mike Walmsley and Chris Lintott and Geron Tobias and Sandor J Kruk and Coleman Krawczyk and Kyle Willett and Steven Bamford and William Keel and Lee S Kelvin and Lucy Fortson and Karen Masters and Vihang Mehta and Brooke Simmons and Rebecca J Smethurst and Elisabeth M L Baeten and Christine Macmillan},
issue = {3},
journal = {Monthly Notices of the Royal Astronomical Society},
month = {12},
pages = {3966-3988},
title = {Galaxy Zoo DECaLS: Detailed Visual Morphology Measurements from Volunteers and Deep Learning for 314,000 Galaxies},
volume = {509},
url = {https://arxiv.org/abs/2102.08414},
year = {2022},
}

You might be interested in reading papers using Zoobot:

• Galaxy Zoo DECaLS (first use at Galaxy Zoo)
• A Comparison of Deep Learning Architectures for Optical Galaxy Morphology Classification
• Practical Galaxy Morphology Tools from Deep Supervised Representation Learning
• Towards Foundation Models for Galaxy Morphology (adding contrastive learning)
• Harnessing the Hubble Space Telescope Archives: A Catalogue of 21,926 Interacting Galaxies

Many other works use Zoobot indirectly via the Galaxy Zoo DECaLS catalog.