Quick Start

This guide walks you through the full MorphoCLIP pipeline: downloading data, extracting features, and training the model.

Command Reference

All commands are shown with both PDM (for local dev) and direct Python (for HPC/conda).

Task	PDM	Python (HPC)
Run tests	`pdm test`	`pytest tests/ -v`
Fetch dataset from S3	`pdm fetch-dataset`	`python scripts/data/fetch_dataset.py`
Extract DINOv3 features	`pdm extract-features`	`python scripts/features/extract_features.py`
Download pre-extracted features	`pdm download-features`	`python scripts/features/download_features.py`
Pre-compute text embeddings	`pdm precompute-text`	`python scripts/text/precompute_text_embeddings.py`
Train model	`pdm train`	`python scripts/training/train.py --config configs/train/base.yaml`
Evaluate model	`pdm eval`	`python scripts/training/eval.py --checkpoint <path>`
Run inference	`pdm infer`	`python scripts/training/infer.py --checkpoint <path>`

Training Pipeline


# 1. Download pre-extracted DINOv3 features from Hugging Face
pdm download-features                              # PDM
python scripts/features/download_features.py       # HPC
 
# 2. Pre-compute text embeddings (once)
pdm precompute-text                                # PDM
python scripts/text/precompute_text_embeddings.py  # HPC
 
# 3. Train MorphoCLIP
pdm train                                          # default config
python scripts/training/train.py --config configs/train/base.yaml
 
# 4. Evaluate a trained checkpoint
python scripts/training/eval.py --checkpoint output/morphoclip_runs/default/checkpoints/best.pt
 
# 5. Run inference (top-k matching, embedding export, or profile export)
python scripts/training/infer.py --checkpoint output/morphoclip_runs/default/checkpoints/best.pt
python scripts/training/infer.py --checkpoint output/morphoclip_runs/default/checkpoints/best.pt --mode embed
python scripts/training/infer.py --checkpoint output/morphoclip_runs/default/checkpoints/best.pt --mode profile

Training Configuration

The default config is at configs/train/base.yaml. Key options:

loss_type: "infonce" (standard contrastive) or "cwcl" (soft-label contrastive)
use_cwa: true to enable Cross-Well Alignment batch correction (disabled by default)
Outputs saved to output/morphoclip_runs/<run_name>/

Alternative configs:

configs/train/mean_pool.yaml — uses mean pooling instead of CrossChannelFormer
configs/train/ddp.yaml — multi-GPU distributed training

Multi-GPU Training

Replace python with torchrun and add --distributed:


# MorphoCLIP on 4 GPUs
torchrun --nproc_per_node=4 scripts/training/train.py --config configs/train/base.yaml --distributed
 
# CellCLIP on 4 GPUs
torchrun --nproc_per_node=4 scripts/cellclip/train_cellclip.py --config configs/cellclip/cellclip_1024dim.yaml --distributed

The batch size in the config is per-GPU, so effective batch = batch_size x num_gpus.

TensorBoard Monitoring

Training logs metrics to TensorBoard automatically:


pip install tensorboard
tensorboard --logdir output/morphoclip_runs/ --port 6006   # MorphoCLIP
tensorboard --logdir output/train_runs/ --port 6006        # CellCLIP

On HPC, forward the port via SSH: ssh -L 6006:localhost:6006 user@hpc-node