TransUNet a 2D Transformer Segmentation Network #2727

Activity

1. finalelement

   added

   Feature request

   on Aug 9, 2021
2. finalelement

   added this to the New transformer arch./layer/blocks milestone on Aug 9, 2021
3. wyli

   added

   Contribution wanted

   on Jan 5, 2023
4. nicolasbeaujoin commented on Oct 22, 2026

   

   nicolasbeaujoin

   on Oct 22, 2026

   More actions

   Hi @finalelement and maintainers,

   I'd like to pick this up if the issue is still current. Quick summary of what I'd propose, so we can align before I open a draft PR:

   Plan
   Add a TransUNet network (from original paper Chen et al., 2021, arXiv:2102.04306) under monai/networks/nets/transunet.py, following the existing structure of UNETR:

   • Hybrid encoder: CNN stem for feature extraction, followed by patch embedding and a stack of monai.networks.blocks.TransformerBlock with positional embeddings.
   • U-Net-like decoder with cascaded upsampling and skip connections from the CNN stem, in the spirit of UnetrUpBlock.
   • API mirroring UNETR (in_channels, out_channels, img_size, hidden_size, mlp_dim, num_heads, num_layers, feature_size, norm_name, ...).

   questions before I start

   1. Backbone: the paper uses a ResNet-50 stem; a lighter 3-conv stem variant would keep the PR self-contained. Is either preferred?
   2. Dimensions: the issue says 2D. is 2D-only acceptable for a first pass, or should I aim for spatial_dims-generic like UNETR, with spatial_dims being either 2 or 3?

   Deliverables

   • monai/networks/nets/transunet.py + export in monai/networks/nets/__init__.py
   • tests/networks/nets/test_transunet.py (parameterized shape tests + TorchScript round-trip, modeled on test_unetr.py)
   • docs/source/networks.rst autoclass entry

   I'll branch off dev, open a draft PR early, and follow the contribution process in CONTRIBUTING.md (formatting via ./runtests.sh --autofix, DCO sign-off).