Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/mwalmsley/zoobot/llms.txt

Use this file to discover all available pages before exploring further.

Zoobot can be installed directly from PyPI for most users, or from source if you plan to modify the library itself. GPU support requires a compatible CUDA installation, but Zoobot also runs correctly on CPU — training is simply slower. Python 3.9 or higher is required.

​
Installation Methods

1

Choose your install method

Pick the method that best fits your workflow:
pip install zoobot[pytorch]
  • PyPI — installs the latest stable release. Best for most users who won’t be modifying Zoobot itself.
  • From source (editable mode) — clones the repository and installs it in editable mode so that local changes to the source are immediately reflected without reinstalling.
  • Google Colab — uses the pytorch_colab extra, which only installs Lightning and avoids reinstalling PyTorch packages that Colab already provides (which can cause conflicts).
Google Colab users must use pip install zoobot[pytorch_colab] rather than pip install zoobot[pytorch]. Reinstalling PyTorch inside Colab can break the runtime environment.
2

(Optional) Set up GPU / CUDA support

If you have an NVIDIA GPU, installing CUDA dramatically speeds up training. Zoobot is tested against CUDA 12.8 with PyTorch 2.7.0.We strongly recommend using conda (or mamba, which is faster) to manage CUDA installation alongside your Python environment:
conda create --name zoobot39_torch python==3.9
conda activate zoobot39_torch
conda install nvidia/label/cuda-12.8.1::cuda
conda install nvidia/label/cuda-12.8.1::cuda-toolkit
Then install Zoobot inside the activated environment:
pip install zoobot[pytorch]
CUDA must be installed before installing PyTorch. If PyTorch is installed without a matching CUDA version, it will default to CPU-only mode. Re-installing into a fresh conda environment is the cleanest fix.
CUDA is entirely optional. Zoobot runs fine on CPU — it is just slower. For experimentation or small datasets, CPU is perfectly adequate.
3

Verify the installation

Confirm that Zoobot and its finetuning module are importable:
import zoobot
from zoobot.pytorch.training import finetune
print('Zoobot installed successfully')
If this runs without errors, you are ready to go. See the Quickstart to finetune your first model.

​
Python and Dependency Requirements

Zoobot requires Python ≥ 3.9. The zoobot[pytorch] extra pulls in the following key dependencies:
PackageMinimum versionPurpose
torch≥ 2.7.0Core deep learning framework
torchvisionlatestImage transforms and datasets
lightning≥ 2.2.5Training loop, checkpointing, early stopping
timm≥ 1.0.15Pretrained encoder architectures (ConvNeXT, MaxViT, EfficientNet, ResNet)
galaxy-datasets≥ 0.0.25CatalogDataModule and galaxy-specific data utilities
huggingface_hublatestAuto-downloading pretrained encoders from HuggingFace
h5pylatestReading HDF5 image files
pandaslatestCatalog management
numpylatestNumerical operations
astropylatestReading FITS image files
scikit-learn≥ 1.0.2Linear sanity-check evaluation
wandblatestOptional experiment tracking (Weights & Biases)
These are installed automatically by pip — you do not need to install them manually.

​
Common Install Issues

Certain CUDA-specific PyTorch builds (e.g. torch==1.12.1+cu113) are not available from pip’s default package index. You need to point pip at the PyTorch package index using the --extra-index-url flag.For example, for CUDA 11.3:
pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
Consult the PyTorch previous versions page to find the exact command for your CUDA version. With Zoobot 2.0 and pip install zoobot[pytorch], this is typically no longer necessary as PyTorch 2.7.0+ supports CUDA 12.8 through the conda method above.
Verify that PyTorch can see your GPU:
import torch
print(torch.cuda.is_available())      # Should print: True
print(torch.cuda.get_device_name(0)) # Should print your GPU model
If torch.cuda.is_available() returns False:
  1. Confirm CUDA is installed: nvcc --version or nvidia-smi in your terminal.
  2. Check that your CUDA version matches what PyTorch expects. CUDA 12.8 is required for PyTorch 2.7.0.
  3. If you installed PyTorch before CUDA, reinstall into a fresh conda environment following the GPU setup steps above.
If you see an import error immediately after installation, the most common causes are:
  • Wrong environment — make sure the environment where you installed Zoobot is the one your Python session is using. Run which python (Linux/macOS) or where python (Windows) to confirm.
  • Editable install not reflecting changes — if you installed from source with -e and are seeing old code, try restarting your Python kernel.
  • Conflicting versions in Colab — use pip install zoobot[pytorch_colab], not the full pytorch extra, to avoid overwriting Colab’s pre-installed packages.
Zoobot 2.0 updated timm to ≥ 1.0.15. Checkpoints saved with older Zoobot versions (which used timm 0.9.x) may not load cleanly with the current version. If you need to load old checkpoints, pin timm to the version used when the checkpoint was saved, or re-export the encoder weights using the old environment before upgrading.
If training crashes with a CUDA out-of-memory error, try reducing the batch size in your CatalogDataModule:
datamodule = CatalogDataModule(
    label_cols=['ring'],
    catalog=labelled_df,
    batch_size=16  # reduce from 32
)
For very limited GPU memory, consider using a smaller encoder (e.g. convnext_nano instead of convnext_base) or setting training_mode='head_only' to freeze the encoder and only train the classification head.

Build docs developers (and LLMs) love

Get started for freeTalk to us