[lufre1]

Closes #187, closes #186.

What changes

The default tiling used to depend only on the VRAM. It now starts from the model's training patch shape and grows by at most 20% when the VRAM allows. Growth is rounded down to multiples of 8 in z and 16 in y/x.

• Model registry: new get_model_training_patch_shape(model_type). Each entry names its source script, and guessed values are marked UNCERTAIN: mitochondria, mitochondria2, cristae–cristae5, vesicles_2d_maus. Please check these.
• torch_em checkpoints: get_prediction reads the patch shape from best.pt (init["train_dataset"]). A UNETR uses exactly its training patch, with no growth. This fixes get_default_tiling is designed to suit UNET models, but not UNETR models. #186: a UNETR resizes each tile to its encoder input, so larger tiles change the scale and can OOM.
• In-memory UNETR without a tiling: raises a ValueError that explains how to pass the tiling.
• Callers: run_segmentation, the CLI and the segmentation widget use the new default. The CLI uses the checkpoint's patch shape when -c is a torch_em checkpoint. The widget resets the tile/halo boxes when the model changes.
• Fixes:
  • get_prediction_torch_em had y/x swapped in block_shape and halo.
  • The widget ignored tiling changes after the first prediction.

Default tiles now

Device	32×256×256 models	64×512×512 (ribbon)	2D models
CPU / 10 GB	32×256×256 (before: 16×96×96 on CPU)	64×512×512	304×304
80 GB	32×304×304 (before: 80×640×640)	72×608×608	304×304 (before: 768×768)

The halo is derived from the patch (min(size // 4, 64)), so it is the same on all devices. Large GPUs now use smaller tiles than before, so inference there will be slower.

Testing

• New tests: TestTiling covers the growth rule with mocked VRAM, 2D and UNETR handling, and reading the patch shape from a checkpoint. A widget test checks that the tiling follows the selected model.
• Full suite (local, CPU): 232/234 pass, 1 skipped. The 2 errors are import failures in test/training because the local torch_em lacks torch_em.model.cryozoo; they are unrelated to this PR.
• End-to-end on tiny trained checkpoints:
  • A supervised U-Net (directory and best.pt) and a mean-teacher checkpoint use exactly their training patch.
  • A plain torch.save model falls back to the VRAM tiling.
  • A 2D UNETR checkpoint uses its 64×64 patch.

Known limitations

• Fixed divisibility: tiles must divide by 8 in z and 16 in y/x. That fits all registered models (depth-4 U-Nets), but a custom U-Net with a different depth may need other values.
• Custom models without stored patch shape: a custom model in the widget, or a plain torch.save file in the CLI, uses the selected model's patch shape.
• Merge conflict: update-backbones adds dtype to the same get_prediction signatures, so a small conflict is expected.

🤖 Generated with Claude Code

[lufre1]

Hi @stmartineau99 ,
could you check if that solves your issue - or what you had in mind?

[stmartineau99]

It works, but I have a few thoughts.

1. I only meant to change the behavior for UNETR models, otherwise use the previous get_default_tiling. Deriving the tile size from the patch_size after loading the model is good. Now the tile sizes for UNETs are much smaller (capped at 1.2x patch size) depending on the VRAM, this would make the prediction slower in those cases. On a related note, I am adding a bfloat16 dtype option that can speed up the predictions by ~2.7x, see 18c2471.
2. I wasn't sure if your "fixed divisibility" comment applies to all models, but it seems it is only a requirement for UNET (based on depth 4, as you said). I assume that there is no similar limitation for UNETR patch size. I set odd values for the z depth, because that's how much could fit in the VRAM; I assume this would not matter here.

[lufre1]

1. That was the point of make tiling dependent on the model's training patch shape #187: almost all of our pretrained U-Nets use InstanceNorm, which normalizes over each tile. A tile much larger than the training patch (e.g. 80×640×640 vs. 32×256×256) feeds the model input statistics it never saw in training. Your bf16 option is the better way to speed things up, since it doesn't change what the model sees. Whichever of our PRs merges second will resolve the small dtype conflict in get_prediction.

2. Right, the divisibility rule only applies when a U-Net tile grows beyond its training patch. A UNETR always gets its exact training patch with no rounding, so odd z values like 25 are kept as they are. I checked that a 25×250×250 patch gives a 25×250×250 tile.

[stmartineau99]

1. That was the point of make tiling dependent on the model's training patch shape #187: almost all of our pretrained U-Nets use InstanceNorm, which normalizes over each tile. A tile much larger than the training patch (e.g. 80×640×640 vs. 32×256×256) feeds the model input statistics it never saw in training. Your bf16 option is the better way to speed things up, since it doesn't change what the model sees. Whichever of our PRs merges second will resolve the small dtype conflict in get_prediction.
2. Right, the divisibility rule only applies when a U-Net tile grows beyond its training patch. A UNETR always gets its exact training patch with no rounding, so odd z values like 25 are kept as they are. I checked that a 25×250×250 patch gives a 25×250×250 tile.

Cool thanks, once this merges, I'll make a PR for the bf16 change.