Code for paper: Low-Memory CNNs Enabling Real-Time Ultrasound Segmentation Towards Mobile Deployment
Train lightweight U-Net models. Code facillitates training of the following models:
- Original U-net as proposed by Ronneberger et al.
- Thin U-Nets with few feature channels per layer (with either regular or separable convolutions)
- Thin, separable convolution U-Nets adapted for knowledge distillation
- Data is expected as two NumPy arrays of grayscale images and corresponding binary labels with dimensions (Batch, Height, Width, Channels = 1)
- No augmentation currently implemented, listed as TODO in
lightweight_unet/dataloader.py kaggle_to_npz.pyscript is included to convert Kaggle data to numpy arrays, while removing all empty images (no salient structure) as detailed in paper
Recommended installation wih conda:
$ conda env create -f environments/environment.yml
Add locations of directories in which Experiment Data and Input Data are to be stored to trainer_cfg in runner.py
Note: Code uses graphviz and pydot to visualize model architectures when saving experiment data and may need separate installation
Main code is run from runner.py, with experiments customizable through model_cfg and trainer_cfg dictionaries. Hard coded modes of operation exist: each experiment must be one of distillation, thin_unet or original_unet. Furthermore, experiments expect a parameter to vary (trainer_cfg['parameter_varied']), which must be selected as a key to one of the configuration dictionaries.
For each experiment, a directory is created (in the path specified in trainer_cfg) with the following contents:
-
Sub-directory for each evaluation fold
- Model architecture as JSON file
- Model weights as .hdf5 file
- Visualisation (.png) of model, containing layer names, shapes etc.
- Training and evaluation information as .csv file. Contains various loss histories and Dice performance of model on evaluation set
-
Log file as .txt detailing configuration of experiments
Finally, a .csv file is also maintained summarising the results of all experiments