# Optimisation of 2D U-Net Model Components for Automatic Prostate Segmentation on MRI
## 3 Materials and Methods
### Model Architecture
From phases 1 and 2, we conclude that the original downsampling and upsampling configurations with a concatenation skip connection in each convolution block, UNet_S2, performs better than the other configurations (see Table 2). Since the phase 2 results correspond with the phase 1 results, where the modified UNet_S is better than the corresponding UNet_S1 and UNet_S2, we perform the next modifications only on the best model. In phase 3, we modify the dropout rate from a value of 0.5 to 0 to confirm the theory discussed in Section 2.6. As can be seen from Table 2 phases 2 and 3, the model with the original dropout rate of 0.5, UNet_S2, still performs better when compared to the modified version, UNet_S2.1, with a dropout rate of 0. In phase 4, we investigate the effect of batch normalisation on the best model, UNet_S2. In model UNet_S2.0.1, the batch normalisation layers are placed in between the convolution and activation layers as shown in Figure 3.  
Figure 2: Skip connection configurations in a convolution block. (**a**) Original convolution block in UNet_S, UNet_S1 and UNet_S2; (**b**) Skip connection with element-wise summation in UNet_S1, UNet_S1.1 and UNet_S2.1; (**c**) Skip connection with concatenation in UNet_S2, UNet_S1.2 and UNet_S2.2. As can be seen from Table 2 phase 4, including batch normalisation layers in model UNet_S.2.0.1 improves on the previous best result, UNet_S.2. The use of the batch normalisation layer also increases the speed of training and reduces training problems related to local minima. As discussed earlier, max pooling is a type of nonlinear filter that preserves the most significant features (lines, edges, etc.) of the previous layer. However, the prostate does not have a well-defined boundary, hence a different pooling layer may perform better segmentation. Therefore, in phase 5, we replace the max pooling layer in the UNet_S.2.0.1 model with average pooling (UNet_S.2.0.1.1), RMS pooling (UNet_S.2.0.1.2) and L2 pooling (UNet_S.2.0.1.3) layers. The model with the average pooling layer, UNet_S.2.0.1.1, is shown to perform better than the UNet_S.2.0.1, UNet_S.2.0.1.2 and UNet_S.2.0.1.3 models as can be observed in Table 2 phase 5. Finally, we investigate options for the activation and batch normalisation layers on the best model, UNet_S.2.0.1.1. As discussed in Section 2.5, the ReLU may cause the neurons in the model to die, which can be overcome by the LReLU. Hence, we replace the ReLU activation layers with LReLU activation layers in UNet_S.2.0.1.1. We also investigate the network performance when applying the batch normalisation layer after each activation layer (as shown in Figure 4) in model UNet_S.2.0.1.1.2. The results are shown in Table 2 phase 6. It can be observed that neither the modification of the ReLU to the LReLU nor the change in position of the batch normalisation layer with the ReLU layer improves on the performance of the UNet_S.2.0.1.1 model. Figure 4: Batch normalisation layer after each activation layer in model UNet_S.2.0.1.1.2. Figure 3: Batch normalisation layer implementations. (**a**) Original convolution and activation layers in model UNet_S.2. (**b**) Batch normalisation layer implementation after each convolution layer in model UNet_S.2.0.1. Therefore, our final model is the UNet_S.2.0.1.1, which is a simplified UNet with a concatenation skip connection in each convolution block, with a \(2\times 2\) average pooling layer used for downsampling and a batch normalisation layer placed between each convolution and ReLU activation layer.