Training Guide

Training Guide

---

Prerequisites

git clone https://github.com/motazalqaoud/Brain-Tumor-Segmentation.git
cd Brain-Tumor-Segmentation
pip install -r requirements.txt

Download the dataset first — see Dataset Setup.

---

Train the 3D Attention U-Net

python scripts/train3d.py \
  --epochs 100 \
  --batch 4 \
  --lr 1e-3 \
  --data-root data/raw/Images_ \
  --base-filters 32 \
  --depth 3

Key Arguments

Argument	Default	Description
--epochs	50	Total training epochs
--batch	4	Batch size
--lr	1e-3	Initial learning rate
--data-root	data/raw/Images_/Images_	Path to dataset
--base-filters	32	Filters in first encoder block (doubles each level)
--depth	3	Encoder/decoder depth
--num-classes	4	Output classes (4 = background + 3 tumor types)
--device	auto	cuda or cpu
--num-workers	0	DataLoader workers (0 = safe for all environments)
--resume	None	Path to checkpoint to resume from

---

Resume Training

Every epoch saves checkpoints/checkpoint_latest.pt. To resume after a shutdown:

python scripts/train3d.py \
  --resume checkpoints/checkpoint_latest.pt \
  --epochs 100

To resume from the best checkpoint:

python scripts/train3d.py \
  --resume checkpoints/best_model_dice_0.XXXX.pt \
  --epochs 100

---

What Gets Saved

File	When	Contains
checkpoints/best_model_dice_X.pt	When val Dice improves	model, optimizer, epoch, metrics
checkpoints/checkpoint_latest.pt	Every epoch	model, optimizer, epoch, metrics
checkpoints/checkpoint_epoch_N.pt	Every 10 epochs	model, epoch, val Dice
visualizations/epoch_NNN_sample_N.png	Every 5 epochs	GT vs prediction comparison
visualizations/training_curves.png	End of training	Loss and Dice curves

---

Train the 2D U-Net (Synthetic Data)

No dataset needed:

python scripts/generate_sample_data.py --n 50 --size 128
python scripts/train.py --epochs 30

---

Hardware Notes

Setup	Recommended config
CPU only	--batch 2 --base-filters 16 --depth 3
8GB GPU	--batch 4 --base-filters 32 --depth 3
16GB+ GPU	--batch 8 --base-filters 32 --depth 4

D_FRAMES = 4 (pseudo-3D depth) can be increased to 8 or 16 on GPUs with more VRAM — edit the constant at the top of scripts/train3d.py.

---

Learning Rate Schedule

ReduceLROnPlateau with factor=0.5, patience=5 — halves the LR when validation loss stops improving for 5 epochs. Logged each time it fires.