docs/custom_data.md

Training on your own data

Pretraining reads WebDataset .tar shards; finetuning is supervised and reads image/mask directories.

data/
  pretrain/                  # stage 1: WebDataset shards
    customdata-000.tar
    customdata-001.tar
    ...
  finetune/                  # stage 2: supervised splits
    train/  images/  masks/
    val/    images/  masks/
    test/   images/  masks/

data/ does not have to be local. Configs support env vars and default reads ${DATA_ROOT:-data}/..., so you can either symlink it

ln -s /mnt/big-disk/affmae-data data

or point the variable elsewhere and leave the repository alone:

export DATA_ROOT=/mnt/big-disk/affmae-data

Relative paths resolve against the repository root, not your shell's working directory, so training scripts work from anywhere. The same applies to CHECKPOINT_ROOT (default weights) and AFFMAE_OUTPUT_DIR (default output).

---

Stage 1 — pretraining data (WebDataset)

Shard layout

Pretraining is masked autoencoding, so just images. Each shard is an uncompressed .tar of images, and webdataset groups files by basename, so one image per basename is all that is required:

customdata-000.tar
├── 000000.png
├── 000001.png
├── 000002.png
└── ...

Three things to consider when building shards:

1. The extension is the key. The loader calls .decode("pil") and then reads sample["png"], so files must be .png. Rename before packing, or change the map_dict(png=...) key in affmae/data/pretrain_dataset.py to match your extension.
2. Shards should be roughly equal in size, ideally 100–1000 images each. Under multi-GPU training each rank gets a disjoint set of shards, so one giant shard leaves other ranks idle.
3. Have at least as many shards as num_workers × ranks. Shards are split across DataLoader workers and across ranks, so a worker that gets none makes WebDataset raise "No samples found in dataset; perhaps you have fewer shards than workers." Three shards with num_workers: 8 fails, and it fails at the first batch rather than at startup.

Building shards

import webdataset as wds
from pathlib import Path

images = sorted(Path("/my/raw/images").glob("*.png"))

# maxcount caps images per shard; %03d must match the brace range in the config
with wds.ShardWriter("data/pretrain/customdata-%03d.tar", maxcount=1000) as sink:
    for i, path in enumerate(images):
        sink.write({"__key__": f"{i:08d}", "png": path.read_bytes()})

ShardWriter prints the shard count when it finishes. If it wrote customdata-000.tar through customdata-031.tar, the config pattern is customdata-{000..031}.tar — the brace range is inclusive and the digit count must match the filenames exactly.

Counting your samples

A WebDataset is an IterableDataset with no length, so the epoch boundary comes from data.total_samples in the config rather than from the data. Getting it wrong does not crash: too high repeats images, too low silently truncates the epoch. Count once:

python -c "
import tarfile, glob
n = sum(sum(1 for m in tarfile.open(f) if m.name.endswith('.png'))
        for f in glob.glob('data/pretrain/*.tar'))
print(n)
"

Running it

cp configs/template_pretrain_512.yaml configs/my_pretrain.yaml
# edit data.path, data.total_samples, data.in_channels
python pretrain.py --config configs/my_pretrain.yaml

Multi-GPU is mostly untested (use at your own rick) with torchrun; data.batch_size is the global batch and is divided across ranks, so the effective batch does not change with GPU count:

torchrun --nproc_per_node=4 pretrain.py --config configs/my_pretrain.yaml

---

Stage 2 — finetuning data (images and masks)

Directory layout

data/finetune/
  train/
    images/   img_0001.tif   img_0002.tif   ...
    masks/    img_0001.tiff  img_0002.tiff  ...
  val/        images/  masks/
  test/       images/  masks/

Things to consider for finetuning on your data:

• Images and masks pair by filename stem. img_0001.tif needs img_0001.tiff. The loader takes the intersection of the two directories, so an unpaired file is skipped without warning.
• Image extension is configurable via data.input_ext (default .tif). The mask extension is fixed at .tiff in affmae/data/finetune_dataset.py.
• val/ is the split evaluated during training. test/ is read only by evaluate.py. All three are required by the templates; point val and test at the same files if you only have two splits.

Mask format

Masks are multi-channel (ie multi-label) TIFFs, one binary channel per annotated structure — not a single-channel label map with integer class IDs.

data.indices selects which channels to train on, and background is added implicitly:

data:
  indices: [0, 1]      # train on mask channels 0 and 1
  num_classes: 3       # len(indices) + 1 for background

So num_classes is typically len(indices) + 1, depending on loss selected, and train.class_weighting must have exactly num_classes entries, background first. A wrong length fails at loss construction with "weight tensor should be defined either for all or no classes".

The selected channels are flattened into a single label map in __getitem__: each channel is binarized at THRESHOLD = 10, and a pixel is assigned position_in_indices + 1. Two consequences:

• Class IDs follow the order of indices, not the channel numbers. indices: [2, 3] makes channel 2 class 1 and channel 3 class 2.
• Channels are applied in order, so on overlapping annotations the later channel wins. Put the class you care most about last.

A single-channel integer label map is not supported — __getitem__ unpacks C, H, W and a 2-D mask raises a ValueError. One-hot it into channels when preparing the data.

Raise a class's weight if it is rare — the shipped [0.2, 2.0, 3.0] deliberately down-weights background.

Running it

cp configs/template_finetune_512.yaml configs/my_finetune.yaml
# edit data.base_path, data.indices, data.num_classes,
#      train.class_weighting, model.pretrained_ckpt_path
python finetune.py --config configs/my_finetune.yaml

You do not have to pretrain first. To start from our released backbone, fetch it once and point the config at the result:

from affmae.data.weights import EMWeights
print(EMWeights.AFFMAE_BASE_PRETRAIN_512.fetch())
# -> weights/pretrain/ckpt_epoch_399_affmae_fpw.pth

model:
  pretrained_ckpt_path: "${CHECKPOINT_ROOT:-weights}/pretrain/ckpt_epoch_399_affmae_fpw.pth"

The encoder architecture keys in your config must match that checkpoint, so keep the aff_* block from the template unchanged unless you are pretraining yourself. See the weights table in the README for every released checkpoint.

---

Checklist before a long run

Cheap to check, expensive to discover at hour six.

Check	Symptom if wrong
data.total_samples matches your shard count	epoch silently truncated or images repeated
Brace range digits match filenames	FileNotFoundError on the first batch
num_classes == len(indices) + 1	crash in the loss, or a dead output channel
len(class_weighting) == num_classes	"weight tensor should be defined either for all or no classes"
in_channels matches your images (1 grey, 3 RGB)	shape mismatch in the patch embedding
Encoder architecture keys identical to pretraining	checkpoint loads with missing/unexpected keys
Dataset size is what you expect	unpaired image/mask stems dropped silently
Shards >= num_workers x ranks	"No samples found in dataset" at the first batch
stats.py matches your data	washed-out or clipped figures; slower convergence
CLAHE removed if not wanted	real intensity differences flattened
Info-bar crop removed for non-EM images	bottom of bright images silently truncated

Sanity-check the pipeline before committing to a full run:

python -c "
from affmae.config import load_config
from affmae.data.finetune_dataset import build_finetune_dataloader
cfg = load_config('configs/my_finetune.yaml')
loader = build_finetune_dataloader(cfg, is_train=True)   # False reads val/
img, mask = next(iter(loader))[:2]
print('image', img.shape, img.dtype)
print('mask ', mask.shape, 'classes present', mask.unique().tolist())
"

mask.unique() should be class indices in [0, num_classes). If you see values above that, indices and num_classes disagree.

---

Adapting the preprocessing

Everything in this section is tuned for greyscale electron micrographs. None of it is wrong for other data, but none of it is neutral either — these are the four places where "trained on EM" is baked in.

1. Images are single-channel

We trained on greyscale. Two consequences:

• data.in_channels: 1 in the config. Set it to 3 for RGB; the patch embedding and the normalization tuple both follow it.
• The pretraining loader converts to greyscale regardless: apply_custom_processing in affmae/data/pretrain_dataset.py calls pil_image.convert("L"). For RGB pretraining, remove that call — otherwise in_channels: 3 silently receives a replicated grey channel.

2. Normalization statistics

affmae/data/stats.py holds four numbers measured on our data:

IMAGE_MEAN = 0.6266           # labelled FPW split: finetuning and inference
IMAGE_STD  = 0.2259
PRETRAIN_IMAGE_MEAN = 0.5562  # unlabelled pretraining corpus
PRETRAIN_IMAGE_STD  = 0.2396

Change these for your own data. They are used by the loaders and by the renderers, which denormalize with the same numbers — so a mismatch shows up as washed-out or clipped figures long before you suspect the statistics.

Measure yours with:

python scripts/calc_norm.py --config configs/my_pretrain.yaml

It reports mean/std over your shards and prints the currently configured values next to them. It measures the raw pipeline deliberately: an earlier version read through the normalized pipeline and reported the residual (~0, ~1), which looks like a passing result and is useless.

For multi-channel data, note these are scalars broadcast across channels ((mean,) * n_channels). Per-channel statistics need create_transforms in affmae/data/pretrain_dataset.py changed to take tuples.

3. CLAHE

Contrast-limited adaptive histogram equalization is applied in four places, because EM micrographs have low local contrast. If your images are already well-exposed, this will flatten real intensity differences and you should remove it:

Where	What	Flag?
affmae/data/preprocess.py	apply_clahe, clipLimit=4.25, 8x8 tiles	yes — preprocess_image(..., use_clahe=False)
affmae/data/finetune_dataset.py:116, 214	same parameters, inline	no flag, edit the call
affmae/data/pretrain_dataset.py:80	clipLimit=4.0, inside apply_custom_processing	no flag, edit the call
affmae/data/pretrain_dataset.py:58	random CLAHE, an augmentation	part of random_transform

Only the inference path takes a flag. The dataset classes hardcode it, so removing CLAHE from training means editing those call sites. If you do, re-measure the normalization statistics afterwards — CLAHE changes the intensity distribution, so the old mean/std no longer describe your inputs.

4. The microscope info bar

apply_custom_processing in affmae/data/pretrain_dataset.py crops the bottom of the image when it finds a row that is more than 95% pure white below the 60% height mark. That removes the scale-bar and metadata strip our microscope burns into every frame.

This is the most EM-specific step in the pipeline and the most likely to damage other data: any image with a genuinely bright bottom edge — a light background, a white border — gets silently truncated, and because the crop is variable-height the resize then changes the aspect ratio per image. Remove it unless your images have the same burned-in strip.

5. Augmentation

ElasticTransform in affmae/data/transforms.py plus affine and photometric transforms are tuned for tissue sections, which deform plausibly. For rigid subjects (documents, industrial parts) elastic warping invents geometry that cannot occur; see scripts/visualize_aug.py to look at what the current pipeline does to your images before training on it:

python scripts/visualize_aug.py --base-path data/finetune --output output/aug.png

Adapting the architecture

Defaults are AFF-Base at 512px with patch size 8. What is worth changing:

Goal	Change
Larger images	data.img_size 768 or 1024, and lower batch_size
Smaller/faster model	shorten aff_depths, narrow aff_embed_dims
More tokens kept per stage	raise aff_ds_rates (0.4 keeps 40% per stage)
Different masking difficulty	model.mask_ratio, pretraining only
Out of memory	lower batch_size, raise train.num_accum to compensate

Two big constraints:

• decoder_depth must equal the number of encoder stages (4 for AFF-Base), one decoder block per stage.
• Encoder keys must match between pretraining and finetuning. The finetune config re-declares aff_embed_dims, aff_depths, aff_num_heads, aff_nbhd_sizes, aff_cluster_size, aff_ds_rates, aff_mlp_ratio and aff_merging_method; any disagreement means the checkpoint does not load into the encoder.

img_size is the exception — it may differ between stages. Cost is driven by the token count, which scales with img_size / patch_size rather than with pixel count, so moving from 512 to 1024 at patch size 8 doubles the grid on each axis.

Leave cluster_attention_backend and decoder_deform_backend at "auto" unless you have profiled the alternative. "auto" selects the fused Triton kernels on CUDA and the PyTorch fallback elsewhere, so the same config runs on CPU, CUDA and Apple silicon; see Hardware and Kernel Coverage in the README.