[nicolasbeaujoin]

Fixes # .

Description

A few sentences describing the changes proposed in this pull request.

Fixes #2727

Adds TransUNet (Chen et al., 2021, arXiv:2102.04306), a hybrid CNN-Transformer segmentation network, following the structure of UNETR:

• Hybrid encoder: a convolutional stem (UnetrBasicBlock) produces skip features at 1/1, 1/2, 1/4, and 1/8 scale, and bottleneck features at 1/16 scale. Each spatial position of the 1/16-scale feature map is tokenized (PatchEmbeddingBlock with patch_size=1, learnable positional embedding) and processed by a stack of TransformerBlocks.
• Decoder: cascaded UnetrUpBlock upsampling from 1/16 back to full resolution, fusing the transformer output with the CNN skip connections.
• API: mirrors UNETR (in_channels, out_channels, img_size, feature_size, hidden_size, mlp_dim, num_heads, num_layers, norm_name, res_block, dropout_rate, qkv_bias, spatial_dims).

Implementation notes / deviations from the paper

• The original is 2D with a ResNet-50 stem. This implementation is dimensionally generic (spatial_dims=2 by default, 2D and 3D supported), consistent with MONAI convention, and uses a lightweight convolutional stem instead of a pretrained R50. A ResNet-based stem and a load_old_state_dict shim for the original checkpoint could follow up if wanted.

Files changed

• monai/networks/nets/transunet.py — network
• monai/networks/nets/__init__.py — export
• tests/networks/nets/test_transunet.py — parameterized 2D/3D shape tests + ill-arg cases
• docs/source/networks.rst — autoclass entry

Testing

• tests/networks/nets/test_transunet.py passes locally (2D and 3D forward shapes, argument validation) against torch 2.5.1.
• CI should confirm the full matrix and the checks this environment lacks (ruff/black/isort, pyrefly).

Draft until review feedback from the maintainers on the open design questions in #2727 is incorporated.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}