Semantic Segmentation (UNet++/U-Net)

The U-Net and UNet++ network models use a similar method for training and evaluation, though they have different network structure settings. Therefore, the U-Net and UNet++ network models are implemented by the same set of code.

In the unet_mindspore and unet_nested_mindspore folders, except for the on_platform/plat_cfg.yaml file, other files are the same.

UNet++/U-Net Training Parameters and Value Ranges

Table 1 describes the names, types, value ranges, default values, and descriptions of the U-Net and UNet++ model parameters.

**Table 1** UNet++/U-Net training parameters (**unet_mindspore**/**model_train.py**)
Parameter	Type	Value Range	Default Value	Description
--device_id	Integer	[0,7]	0	ID of the NPU used for training.
--device_num	Integer	Currently, only one device is supported, that is, the value can only be 1.	1	Number of NPUs used for training.
--epoch_size	Integer	[1,10000]	200	Number of training epochs.
--batch_size	Integer	Less than or equal to the number of images in the dataset. [1, 512]	16	Batch size for training.
--pretrained_ckpt_path	String	-	""	Path for storing the pre-trained model. "" indicates that the pre-trained model is not loaded.
--input_width	String	[1,960], which is a multiple of 16.	"256"	Model input width.
--input_height	String	[1,960], which is a multiple of 16.	"192"	Model input height.
--init_lr	Float	(0,1)	0.0003	Training learning rate.
--train_dataset_path	String	-	None	Path of the training dataset.
--train_output_path	String	-	"train_output_path"	Output path of the training result.
--run_eval	Bool	True or False	True	Whether to perform validation during training.
--model	String	unet or unet_nested	unet	Model name. unet indicates that the U-Net model is selected, and unet_nested indicates that the UNet++ model is selected.
--split_ratio	Float	(0,1)	0.8	Dataset splitting ratio during training and evaluation.

UNet++/U-Net Training Command Reference

Run the following commands to start UNet++/U-Net model training:

The U-Net supports two modes: pretrain ckpt and train from scratch. By default, train from scratch is used. That is, no pre-trained model is loaded, which is recommended. If you need to fine-tune the model, use --pretrained_ckpt_path to add the path of the trained model to the training script.

python3 model_train.py --train_dataset_path={train_dataset_path} --train_output_path=./unetpp_output_dir --epoch_size=50 --batch_size=8 --input_width=256 --input_height=192 --init_lr=0.0005 --device_num=1 --device_id=1 --model=unet_nested --run_eval=True

The model training process is random. Use the evaluated accuracy. Figure 1 shows the training accuracy results.

Figure 1 UNet++/U-Net training accuracy results

Figure 2 shows the log information after the model training is complete.

Figure 2 Information about the UNet++/U-Net training

After the model training is complete, the .ckpt, .a310.om, .a310p.om, and .air model files are generated in the output directory specified by the --train_output_path parameter.

UNet++/U-Net Evaluation Parameters and Value Ranges

Table 2 describes the names, types, value ranges, default values, and descriptions of the evaluation parameters of the UNet++/U-Net model.

**Table 2** UNet++/U-Net evaluation parameters (**unet_mindspore**/**model_eval.py**)
Parameter	Type	Value Range	Default Value	Description
--eval_dataset_path	String	-	None	Path of the evaluation dataset.
--eval_ckpt_path	String	-	None	Path for obtaining the evaluation CKPT. You need to specify the training output path.
--device_id	Integer	[0,7]	0	ID of the NPU used for training.
--eval_output_path	String	-	"eval_output_path"	Output path of the evaluation results.
--model	String	unet or unet_nested	"unet"	Model name. unet indicates that the U-Net model is used, and unet_nested indicates that the UNet++ model is used.

UNet++/U-Net Evaluation Command Reference

Run the following command to start UNet++/U-Net model evaluation:

python3 model_eval.py --eval_dataset_path={eval_dataset_path} --eval_ckpt_path=./output_dir --eval_output_path=./eval_result --device_id=0 --model=unet_nested

Use the checkpoint generated during the training to evaluate the model accuracy, as shown in Figure 3.

Figure 3 UNet++/U-Net evaluation results

The folders shown in Figure 4 are generated in the evaluation directory. The all_images folder stores the evaluation result of each image, and the every_images_statistics.csv folder stores the IoU value and Dice coefficient of each image, and the statistics.csv folder stores the average IoU value and Dice coefficient.

Figure 4 UNet++/U-Net evaluation directory

Parent topic: Training Component