Implementation of the U-Net Transformer and comparison with Attention U-Net and baseline U-Net

• Quick Start
• Description
• Training & Hardware
• Results
• Requirements
• Data

1. Create Environment

conda create -n <environment_name> python=3.8
conda activate <enviroment_name>
pip install -r requirements.txt

2. Download dataset as mentioned in Data section
3. Run the following commands only one time to place data (If you have already run this command once make the is_raw flag to False):

train_ratio = 0.65
val_ratio = 0.20
batch_size = 256
size = (128, 128)
num_workers = 2
is_raw = True

train_loader, val_loader, test_loader = get_loaders(f'../data/{"raw" if is_raw else "processed"}/lgg-mri-segmentation/kaggle_3m', train_ratio, val_ratio, batch_size, size, num_workers, is_raw)
print(f'Train samples: {len(train_loader.dataset)}')
print(f'Val samples: {len(val_loader.dataset)}')
print(f'Test samples: {len(test_loader.dataset)}')

4. These commands are written in provided notebooks in notebooks/ directory
5. Solver.py class trains and evaluate model(s)

This model was trained on Brain MRI Segmentation Kaggle dataset and scored a Dice Score of 0.8925 on validation data and 0.8868 on test data.

This model can easily be extended for multiclass classification. Moreover, I wrote U-Net generlized implementation which can be easily be extended for different types of U-Nets.

• Attention UNet was run on 2 GPUs such that batch was divided in half (128 on each)
• UTransformer was run on 8 GPUs such that I had to run certain layers manually on each GPU:
  • Input, output, loss, and PE on cuda:0
  • Encoder layers on cuda:6
  • Decoder layers on cuda:7
  • MHSA on cuda:1
  • Each (4)MHCA on cuda:2 through cuda:5

1. Anaconda
2. CUDA 11.3 or later
3. PyTorch 1.12 or later
4. Jupyter Notebook
5. 8 GPUs to run U-Net Transformer

Brain MRI Segmentation can be downloaded from Kaggle.

• Place the lgg-mri-segmentation folder with its contents into data/raw folder
• Run ETL.py

Olivier Petit, Nicolas Thome, Clément Rambour, Luc Soler:

U-Net Transformer: Self and Cross Attention for Medical Image Segmentation