UNet3D MLPerf

[francislata]

Overview

This is the training script for the UNet3D MLPerf model.

How to run training script

Download dataset

Run the following script from the root folder of tinygrad:

./examples/mlperf/scripts/setup_kits19_dataset.sh

Training model

For training on a tinybox green:

time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py