Skip to content

Latest commit

 

History

35 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Step by step for running a model

Install repository

pip install -e  git+https://github.com/LOFAR-VLBI/astroNNomy.git#egg=astroNNomy

DINOv3 additionally needs the third_party/dinov3 submodule, which pip install does not fetch. Clone with git clone --recurse-submodules, or in an existing checkout:

git submodule update --init third_party/dinov3

Get training data

export FITS_DIR = ...
mkdir $FITS_DIR
cd $FITS_DIR

ONLINE_DATA_PATH=https://public.spider.surfsara.nl/project/lofarvwf/jdejong/CORTEX/calibrator_selection_robertjan/cnn_data/
wget -r --spider --no-parent $ONLINE_DATA_PATH 2>&1 | grep -o 'https://[^ ]*' | grep -E '^.*\/{1}[^\/]+?\.[^\/]+?$' > urls.txt
xargs -n 1 -P 16 wget -q -r -np --no-clobber < urls.txt 

Preprocess fits files and convert to npy

export DATA_DIR="${FITS_DIR}_1024"
python -m astronnomy.pre_processing_for_ml $FITS_DIR --out-dir $DATA_DIR --resize 1024 --print-stats

This writes uncompressed *.npy arrays plus a stats.json sidecar, mirroring the class/mode folder layout of the input. Options:

flag default meaning
--out-dir root_dir Where to write. Defaults to alongside the *.fits, which requires a writable source dir.
--resize 2048 Store images at this size, downsampled with an antialiasing filter. 2048 is the native size, i.e. no resampling.
--dtype float16 Storage dtype. The values live in [0, 1], so float16 costs ~2e-4 of precision.
--max-scale 0.2 Skip images with a pixel scale of at least this many arcsec/pixel.

Resize once, here, not in the dataloader. The network never sees more than data_transforms.resize_max pixels (896 by default), and RotateAndCrop takes a size/sqrt(2) crop before that, so storing at 2048 means every worker pays to load and resample data that is immediately thrown away. At --resize 1024 an image is 2.1 MB instead of 14.8 MB and __getitem__ costs ~2 ms instead of ~160 ms. Going below ~1270 starts upsampling at the top of the resize_minresize_max range.

Dataset statistics (mean/std) are accumulated during the same pass and written to $DATA_DIR/stats.json, so training reads them instead of making its own pass over the data. Older *.npz directories without a sidecar still work: the statistics are then computed on first use and cached (see below).

(Optional) Copy files to /dev/shm for fast dataloading

find $DATA_DIR -type f -name "*.npy" | xargs -n 1 -P 8 -i rsync -R {} /dev/shm
rsync -R $DATA_DIR/stats.json /dev/shm
DATA_DIR = /dev/shm$DATA_DIR

At --resize 1024 the full ~15k-image set is ~32 GB, which fits in /dev/shm directly. The squashfs mount that used to be needed for the 223 GB of *.npz is no longer necessary — and it cost a second decompression layer on every read.

Run neural network training

python -m astronnomy.training.train_nn src/astronnomy/training/config.yaml \
    dataloader.dataset_root=/dev/shm$DATA_DIR

The first argument is a config file; any number of section.key=value overrides may follow. The dataloader expects the filetree to be in the following format:

<filepath>
  |- continue
  |- stop
  |- continue_val
  |- stop_val
  |- stats.json     (written by the preprocessing step)

Dataset statistics cache

When there is no stats.json (e.g. an older *.npz directory), the statistics are computed once and memoized. The cache lives outside the dataset root, since that is often read-only:

  1. dataloader.cache_dir in the config, if set,
  2. else $ASTRONNOMY_CACHE_DIR,
  3. else $XDG_CACHE_HOME/astronnomy/dataset_stats (i.e. ~/.cache/astronnomy/dataset_stats).

The cache key is derived from the file names rather than their absolute paths, so a cache built against the shared filesystem still applies after the data is staged to /dev/shm.

Augmentation noise

data_transforms.noise_sigma (default 0.0) adds Gaussian noise of that absolute sigma after resampling, as an SNR-jitter augmentation. Resampling itself is antialiased, so noise is not needed to cover it. Values around 0.020.05 are a reasonable starting point against a typical per-image std of ~0.16.

Run neural network training with DINOv3

DINOv3 needs two things DINOv2 does not: its source has to be a local checkout (that is the third_party/dinov3 submodule, see the install section), and its weights are access-gated, so they cannot be fetched at runtime.

  1. Request access at https://ai.meta.com/resources/models-and-libraries/dinov3-downloads/ — you get an email with download URLs.
  2. wget the checkpoint for the variant you want into a directory of your choice, keeping the upstream filename (e.g. dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth).
  3. Point the dinov3.weights key in the config at that directory (a direct .pth path also works):
dinov3:
  weights: /projects/prjs1845/dinov3_model_weights

The variant is then selected like any other backbone:

python -m astronnomy.training.train_nn src/astronnomy/training/config.yaml \
    model.model_name=dinov3_vitb16 \
    data_transforms.resize_min=640 data_transforms.resize_max=640 data_transforms.resize_val=640

The filename is matched as <model_name>_*.pth inside the directory, so the hash suffix never has to be written down and several variants can live side by side.

Two environment variables override the defaults, for machines where the config's paths don't hold:

variable overrides
DINOV3_WEIGHTS the dinov3.weights config key. Takes over whenever the configured path does not exist, which is what lets a checkpoint trained elsewhere be reloaded here.
DINOV3_REPO_DIR the third_party/dinov3 submodule, for a non-editable install or a clone you are modifying.

Supported DINOv3 models

Names are the upstream hubconf.py entrypoints:

  • dinov3_vits16 (21M params)
  • dinov3_vits16plus (29M params)
  • dinov3_vitb16 (86M params)
  • dinov3_vitl16 (300M params)
  • dinov3_vitl16plus (~330M params)
  • dinov3_vith16plus (840M params)
  • dinov3_vit7b16 (6.7B params)

Important: DINOv3 uses patch size 16 (vs DINOv2's 14), so image resizing must be a multiple of 16. It also uses RoPE instead of a learned positional embedding, so model.pos_embed is ignored, and its register tokens are called storage tokens — model.tune_register_tokens drives both.

Acknowledgments

This repository is part of the project CORTEX (NWA.1160.18.316) of the research programme NWA-ORC which is (partly) financed by the Dutch Research Council (NWO).

About

Neural Networks for LOFAR

Resources

Stars

0 stars

Watchers

3 watching

Forks

Releases

Packages

Used by

Contributors

Languages