Spaces:

smerkd
/

affmae

Running on Zero

App Files Files Community

affmae / docs /custom_data.md

smerkd

AFF-MAE ZeroGPU demo: segmentation, reconstruction, token layout, batch

8065faa verified 2 months ago

preview code

Raw History Blame Contribute Delete

14.2 kB

# 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

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

or point the variable elsewhere and leave the repository alone:

```bash
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
```python
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:
```bash
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
```bash
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:

```bash
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:

```yaml
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

```bash
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:

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

```yaml
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:

```bash
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:

```python
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:

```bash
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:

```bash
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.

	# 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

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

	or point the variable elsewhere and leave the repository alone:

	```bash
	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

	```python
	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:

	```bash
	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

	```bash
	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:

	```bash
	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:

	```yaml
	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

	```bash
	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:

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

	```yaml
	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:

	```bash
	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:

	```python
	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:

	```bash
	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:

	```bash
	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.