Practical overview
When it fits
Use X Unet when you need a configurable PyTorch U-Net implementation and want to experiment with nested stages or learned cross-stage attention, including for 3D-shaped inputs.
Who it is for
- Python developers building or experimenting with U-Net-based image models.
- Researchers and engineers exploring learned cross-stage attention or nested U-Net configurations.
- Developers prototyping models for video or volumetric medical-image inputs.
Documented uses
- Construct a 2D U-Net model and pass image-shaped tensors through it.
- Configure nested stages and consolidated upsampling feature maps.
- Replace standard skip connections with learned cross-stage attention residuals.
- Set up a 3D model for video or CT/MRI-shaped tensors.
- Use XAttnRes separately to route features across multiple scales.
Requirements & limits
The public repository provides installation instructions using pip and specifies an MIT license.
- The examples show model construction and tensor input/output shapes rather than a complete training or evaluation pipeline.
- Cross-stage attention can be enabled alongside skip connections or configured to replace them.
- The repository includes 3D video and CT/MRI-shaped examples, while listing 3D memory-efficiency work as a todo.
Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.
Install X Unet and construct a 2D model
Install the library and create the documented 2D XUnet model with a sample tensor.
Before you start
- Python environment with PyTorch installed
- Package installation via pip
Steps
Install the package with the command below.
pip install x-unetImport
torchandXUnetfromx_unet.Construct
XUnetwithdim=64,channels=3,dim_mults=(1, 2, 4, 8),nested_unet_depths=(7, 4, 2, 1),consolidate_upsample_fmaps=True, anduse_xattnres=True.Create a sample tensor with
torch.randn(1, 3, 256, 256)and pass it to the model.
Expected outcome
The example produces an output tensor with shape (1, 3, 256, 256).
Guide limitations
- The example demonstrates model construction and a tensor pass; it does not document training, dataset preparation, or evaluation.
Construct the documented 3D X Unet model
Set up the repository's example model for video or CT/MRI-shaped inputs.
Before you start
- Python environment with PyTorch installed
- Package installed with
pip install x-unet
Steps
Import
torchandXUnetfromx_unet.Construct
XUnetwithdim=64,frame_kernel_size=3,channels=3,dim_mults=(1, 2, 4, 8),nested_unet_depths=(5, 4, 2, 1),consolidate_upsample_fmaps=True,weight_standardize=True, anduse_xattnres=True.Create a sample tensor shaped
(1, 3, 10, 128, 128)and pass it to the model.
Expected outcome
The example specifies an output with the same shape as its 3D input.
Guide limitations
- The repository lists memory-efficiency improvements for 3D as todo items; the example does not establish performance or memory requirements.
Sources & review
Source record details
- Source availability
- Active
- Sources checked
- 8 October 2026
- Original website title
- lucidrains/x-unet · Official homepage title · retained in its source language

