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 50 \
  --batch 4 \
  --lr 1e-3 \
  --data-root data/raw/Images_ \
  --base-filters 32 \
  --depth 2

Or use a preset config (recommended):

python scripts/train3d.py --config configs/cpu.json     --data-root data/raw/Images_
python scripts/train3d.py --config configs/gpu_8gb.json --data-root data/raw/Images_

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_	Path to dataset
--base-filters	32	Filters in first encoder block (doubles each level)
--depth	2	Encoder/decoder depth
--num-classes	8	Output classes (8 = background + 7 WHO tumor categories)
--device	auto	cuda or cpu
--num-workers	0	DataLoader workers (0 = safe for all environments)
--resume	None	Path to checkpoint to resume from
--image-size	64	Resize spatial dimensions before training
--d-frames	2	Pseudo-3D depth (slices stacked per sample)

---

Resume Training

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

python scripts/train3d.py \
  --config configs/cpu.json \
  --data-root data/raw/Images_ \
  --resume checkpoints/checkpoint_latest.pt

The checkpoint restores model weights, optimizer state, and learning rate scheduler state exactly.

---

What Gets Saved

File	When	Contains
checkpoints/best_model_dice_X.pt	When val Dice improves	model, optimizer, scheduler, epoch, metrics
checkpoints/checkpoint_latest.pt	Every epoch	model, optimizer, scheduler, 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

---

Training Output

Each epoch logs:

Train Loss: X.XXXX | Mean Tumor Dice: X.XXXX
  Glioma: X.XXXX | Meningioma: X.XXXX | Nerve Sheath: X.XXXX | Embryonic: X.XXXX | Mixed Neuronal: X.XXXX | Mesenchymal: X.XXXX | Germ Cell: X.XXXX
Val   Loss: X.XXXX | Mean Tumor Dice: X.XXXX
  Glioma: X.XXXX | Meningioma: X.XXXX | ... | Background: X.XXXX

After epoch 50, a full test set evaluation runs automatically:

TEST SET RESULTS
Mean Tumor Dice: X.XXXX
  Glioma: X.XXXX | Meningioma: X.XXXX | Nerve Sheath: X.XXXX | Embryonic: X.XXXX | Mixed Neuronal: X.XXXX | Mesenchymal: X.XXXX | Germ Cell: X.XXXX
Background Dice: X.XXXX
Test Loss: X.XXXX

---

Hardware Notes

Config	Hardware	Approx. time/epoch	Settings
configs/cpu.json	CPU	~43 min (measured)	depth=2, image=64, D=2
configs/gpu_8gb.json	RTX 3070 / 4060 Ti	~2 min (estimated)	depth=3, image=128, D=4
configs/gpu_16gb.json	RTX 3090 / 4090 / A100	~45 sec (estimated)	depth=4, image=128, D=8

---

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.