BraTS Brain Tumor Segmentation — TransUNet and U-Net checkpoints
Weights for a controlled comparison of U-Net and TransUNet on BraTS-style multi-modal brain MRI, trained from scratch under a fixed 15-epoch budget with identical data and optimization settings across model families.
Code and full write-up: https://github.com/yhj3/brats-unet-vs-transunet
Which file to use
| File | Model | FG Dice | mean CE | Dice 1 / 2 / 3 |
|---|---|---|---|---|
checkpoint_jobA_epoch19.pth |
TransUNet baseline (R50-ViT-B/16), 320² | 0.7578 | 0.0464 | 0.6775 / 0.7661 / 0.8299 |
checkpoint_improved_transunet_epoch20.pth |
A separate TransUNet run | — | — | — |
Use checkpoint_jobA_epoch19.pth. It is the best checkpoint and the one reported in the
paper table and on the project page. Despite its name, checkpoint_improved_transunet_epoch20.pth
is from a different run and is not the reported model; it is kept only for reference.
Context
Foreground Dice is averaged over labels 1–3 and computed only over slices where the class is present — including absent classes inflates the score, since a model then gets credit for predicting "nothing here" on the majority of slices. Numbers here are not comparable to Dice computed over all classes.
For reference, on the same split and metric:
| Method | Input res. | mean CE ↓ | FG Dice ↑ |
|---|---|---|---|
| Original U-Net | crop | 0.0437 | 0.7326 |
| Flex U-Net (norm + depth only) | crop | 0.3850 | 0.6323 |
| Flex U-Net (loss improved) | crop | 0.1952 | 0.7033 |
| TransUNet baseline | 320² | 0.0455 | 0.7575 |
| TransUNet (skip refinement + aug) | 224² | 0.0411 | 0.7347 |
TransUNet's advantage is not uniform: on the compact structures the tuned U-Net is competitive or better (Dice 3: 0.8351 vs. 0.8322), and essentially all of the gain sits in label 2 — the diffuse subregion — at 0.7652 vs. 0.6800.
Input format
4-channel 2D axial slices — FLAIR, T1, T1-CE, T2 — z-score normalized per slice, with a brain crop computed from non-zero tissue. Four output classes (BraTS label 4 is remapped to 3). The checkpoint stores the model configuration and the crop coordinates alongside the weights, so evaluation reconstructs the exact model:
git clone https://github.com/yhj3/brats-unet-vs-transunet && cd brats-unet-vs-transunet
pip install -r requirements.txt
hf download yihangj3/brats-transunet checkpoint_jobA_epoch19.pth --local-dir checkpoints/
python src/eval.py --checkpoint checkpoints/checkpoint_jobA_epoch19.pth
Training data
BraTS multi-modal MRI (FLAIR, T1, T1-CE, T2). The dataset is not redistributed here; obtain it from the BraTS challenge organizers.
Attribution
The TransUNet encoder/decoder modules are adapted from the official TransUNet implementation (Chen et al., 2021), MIT licensed. The 4→3 channel projection, patch-grid handling at non-default input sizes, and the decoder-side skip refinement are my additions.
Contact
Yihang Jiao — yihangj3@illinois.edu · https://yhj3.github.io