To generate the memory consumption model for 3D U-net, we first begin by collecting performance samples by running the 3D U-net model training for various image dimensions and batch sizes. The memory profiler mprof was used to profile the memory consumption during the training phase, and the peak memory consumed was extracted during execution (Step 1 in Figure 6). Intel’s optimized Tensorflow version 2.4 was used, and the application was written in Keras (version 2.2.3). More details about the software stack and the hardware utilized to collect the performance samples are described in the reproducibility section in the Appendix.
To facilitate the data collection for various input dimensions and batch sizes, we have generated scripts that create random input data for the desired image dimensions. The synthetic input data is used for the model training. The memory consumption was sampled throughout execution using the mprof profiler (Step 2 in Figure 6). Mprof polls the program to collect the memory consumption details at a sampling interval (default is 0.1 seconds). When the training is complete, we extract the peak memory consumed as our performance sample for that run. Table 1 shows the parameter space used for collecting performance samples.
Table 1. Parameter design space used for memory consumption modeling for 3D U-net
Parameter | Range |
Image dimension (along each dimension) | 32,64,128,144,200,244 |
Batch size | 1,2,4,8,16,32,64 |
When all the performance samples have been collected, we randomly classify the data into training data and test data (Step 2). 70% of the total collected samples was used for training and the remaining 30% for test/validation data. Random sampling is used to eliminate any bias in training data and to help generate a model that can generalize well to the test data.
The training data was inputted into the Pandora symbolic regression tool to generate the memory consumption model (Step 3 in Figure 6). The performance metric under consideration is the peak memory consumption during execution, and the input parameters are the image dimension and batch size of the 3D U-net. The configuration parameters used for model generation are those discussed in Model generation process.