A Pruning Method for U-Net Network Model Based on Genetic Algorithm

Abstract

The present invention discloses a pruning method for the U-Net network model based on the genetic algorithm. First, the present invention performs amplification processing on the data set. Secondly, the genetic optimization algorithm is combined with the U-Net model, and the loss function of the U-Net model is introduced into the fitness index of the genetic algorithm, and the optimization target is the compression rate combination of the channel numbers of each layer of the U-Net model. Subsequently, the compression rate combination of the required optimal U-Net network is obtained through the genetic algorithm, and pruning processing is performed on the network structure. Compared with the U-Net network model before pruning, on the segmentation results of three data sets and their respective amplified data sets, there is no significant difference in the results before and after model pruning, but the number of parameters and the amount of calculation of the entire U-Net network are significantly reduced, achieving the effect of network model pruning.

Description

Technical Field

[0001] The present invention relates to a pruning method for a U-Net network model, specifically a method that combines the U-Net network in deep learning with a genetic algorithm for pruning. This method is actually applied to the segmentation of retinal vascular images. Background Art

[0002] The extraction of retinal blood vessels and the characterization of morphological attributes, such as diameter, shape, tortuosity, and bifurcation, can be used for screening, evaluating, and treating different ocular abnormalities. With the rapid development of artificial intelligence technology in recent years, neural network models have been widely applied in the field of image segmentation. However, the current mainstream neural network models all have very complex network structures, so the training time of the models needs to be greatly increased. As the network models become deeper and deeper, while the network performance is improved, another problem arises: the more complex network structures bring greater consumption of computing resources, and unreasonable settings in the design of the number of convolutional kernels will lead to insufficient utilization of convolutional layer parameters, ultimately wasting a large amount of computing resources. In recent years, the U-Net network model used in the field of medical image segmentation has good segmentation effects due to its unique structure. However, for the setting of the number of channels in each layer of the U-Net, it often depends on the manual setting of researchers, which makes the segmentation effect of the network rely to a large extent on manual experience. In view of this phenomenon, in the work of segmenting retinal blood vessels, this paper proposes a method for automatically pruning the U-Net channels using a genetic algorithm, aiming to obtain a more lightweight U-Net network structure through channel pruning, with performance no less than that of the original structure. Summary of the Invention

[0003] Aiming at the problem of parameter redundancy existing in the existing U-Net network model in retinal image segmentation, the present invention proposes a pruning method for the U-Net network model based on a genetic algorithm. Different from most previous channel pruning methods, what is pruned is the number of channels in each layer of the U-Net network structure, rather than selectively pruning "unimportant" channels. By setting a specific compression rate for each layer and using the fast randomness of the genetic algorithm, a combination of optimal compression rates with better robustness can be obtained as much as possible.

[0004] According to the technical solution provided by the present invention, a U-Net pruning method based on a genetic algorithm is proposed, including the following steps:

[0005] Step 1: Horizontally flip, vertically flip, and rotate the training set images at multiple angles;

[0006] Step 2: Perform preprocessing operations on the color images;

[0007] Step 3: Construct a new loss function for the U-Net network model to make it more suitable for the training of retinal image datasets with pixel imbalance problems;

[0008] The expression of the new loss function is as follows:

[0009] Loss = L dice + λL r (1)

[0010] L r is the cross-entropy function, and the expression is as follows:

[0011]

[0012] where TP and TN are the numbers of true positive and true negative pixels respectively; N p and N n are the numbers of target pixels and non-target pixels respectively; y is the label value, y = 1 for the segmentation target, y = 0 for the background; p is the predicted probability value of the pixel;

[0013] L dice is the Dice coefficient expression, as follows:

[0014]

[0015] where N is the number value of pixels; p(k,i) ∈ [0,1], q(k,i) ∈ [0,1] are the predicted probability and true label of pixel point k class respectively, and λ is the coefficient;

[0016] Step 4: Take the number of four symmetric convolutional layers of the U-Net network model as the optimization target, that is, the optimization result of the genetic algorithm is a combination containing four compression rate values; use the genetic optimization algorithm to optimize this combination, and finally obtain an optimal compression rate combination. After performing corresponding modification operations on the number of each convolutional layer, a pruned U-Net network model is obtained;

[0017] Step 5: Incorporate the loss function of the U-Net network into the fitness function of the genetic algorithm to make them have an inverse relationship, and expand the gap between the loss functions through the fitness function;

[0018] Step 6: Through continuous iteration of the genetic algorithm, obtain the required optimal compression rate combination, and conduct experimental comparisons before and after pruning the U-Net network model on three datasets respectively.

[0019] Preferably, perform horizontal flipping and vertical flipping on the training set images to increase the data volume by 4 times; perform multi-angle rotation on the training set images to increase the data volume by 8 times.

[0020] Preferably, in the step 4, the compression ratio combinations are as follows:

[0021] α = [α 1 , α 2 , α 3 , α 4 (4)

[0022] wherein the value ranges of α 1 , α 2 , α 3 , α 4 are [0.1, 0.2, …, 0.8]; α is an input parameter of the genetic algorithm.

[0023] Preferably, the U-Net network structure is as follows:

[0024] In the encoding structure of the network, four encoding blocks are used, and each encoding block includes two convolutional layers, a batch normalization layer, an LReLU activation layer, and a max pooling layer;

[0025] In the decoding structure of the network, corresponding to the encoding structure, four decoding blocks are used, and each decoding block includes two convolutional layers, an upsampling layer, a BN layer, and a skip connection layer.

[0026] Preferably, in the step 5, the expression of the fitness function in the genetic algorithm is:

[0027] Fitness = e 20*(0.3-Loss) (5)

[0028] where Loss is the loss function of the U-Net network model.

[0029] Preferably, the pseudocode for the combination of the genetic algorithm and the U-Net network model is as follows:

[0030]

[0031]

[0032] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0033] Compared with the original structure of the U-Net network model, the method of the present invention performs corresponding pruning on the model structure. While reducing the number of network parameters, it does not cause a significant decrease in the segmentation result, achieving a good pruning expectation; compared with the traditional method of manually setting the number of channels in each layer of the U-Net, the method of the present invention has the advantage of automatically optimizing to obtain the best number of channels, and after a series of optimization operations of the genetic algorithm, the obtained results are more persuasive. BRIEF DESCRIPTION OF THE DRAWINGS

[0034] Figure 1 is the flowchart of the present invention;

[0035] Figure 2 is the structural diagram of the U-Net network model of the present invention;

[0036] Figure 3 is the schematic diagram of convolutional layer pruning;

[0037] Figure 4 is the graph of the change of fitness function value;

[0038] Figure 5 is the comparison graph of the results of the U-Net before and after pruning on the DRIVE augmented dataset. Detailed implementation manners

[0039] The present invention will be further described below in conjunction with specific embodiments. The following description is only for demonstration and acceptance, and does not impose any formal restrictions on the present invention.

[0040] As Figure 1 shown, the implementation steps of the embodiments of the present invention are as follows:

[0041] Step 1: Perform data augmentation operations on the training sets of the existing public datasets DRIVE, STARE, and CHASE_DB1. Specifically, perform horizontal flipping, vertical flipping, and multi-angle rotation on the images, and amplify the data volume to 4 times and 8 times the original respectively. The purpose of this step is on the one hand to provide more training sets for model training, and on the other hand to verify whether the U-Net network model can still maintain the stability of the segmentation results before and after pruning when using different amounts of training sets for training.

[0042] Step 2: The preprocessing process of the images is as follows:

[0043] Separate the channels of the color images, select the green channel with better blood vessel clarity as the input image for processing, and the image size is 576×576;

[0044] Step 3: Construct a new network model loss function for the U-Net in order to better feedback the training process during training, so that the model can obtain better training parameters. For this reason, the method of the present invention takes into account the characteristics of pixel distribution imbalance existing in all three retinal image training sets used. On the basis of the traditional binary cross-entropy loss function, the Dice coefficient is introduced, and the two are weighted and combined to form a new loss function. The specific expression is shown in Equation (1), and the U-Net network model structure is as Figure 2 shown.

[0045] Step 4: Perform symmetric pruning on the symmetric structure of the U-Net network model. The target of pruning is the convolutional layers with a large number of parameters in each layer, and reducing the number of convolutional layers is the key point of pruning. On the basis of the original structure quantity, a compression rate parameter is set for each layer, and a compression rate combination containing four parameters is obtained. Subsequently, a genetic optimization algorithm is used to optimize this combination. Finally, an optimal compression rate combination is obtained. After performing corresponding modification operations on the number of convolutional layers in each layer, a pruned U-Net network model is obtained. The pruning operation of the convolutional layer is as Figure 3 shown.

[0046] Step 5: In order to evaluate the pruning effects produced by different compression rate combinations, the minimum loss function obtained from training the U-Net network model is used as a parameter and input into the fitness function of the genetic algorithm, so that there is an inverse relationship between the two, and the fitness function is used to expand the gap between the loss functions, so as to facilitate better iterative calculation of the next generation in the genetic algorithm. Regarding the relevant parameter settings of the U-Net model, the optimizer uses the stochastic gradient descent algorithm (SGD) to optimize the parameter training of the model. The initial learning rate is set to 0.1, and it is multiplied by 0.1 every twenty epochs as the training iteration times decrease, and the epoch is set to 60. In order to accelerate the optimization speed of the genetic algorithm, the input image is sliced into 48×48 pixel blocks for training.

[0047] Step 6: For the genetic algorithm, the mutation rate and recombination probability are set to 0.01 and 0.8 respectively, the population size is set to 100, and the number of iterations is 70. Experiments show that the compression rate combination has reached the optimal solution of the algorithm after 45 generations of iteration and no longer updates. The specific fitness change is as Figure 4 shown.

[0048] Table 1 Comparison of the U-Net network model before and after pruning

[0049]

[0050]

[0051] Table 1 lists the comparison of relevant indicators of the U-Net network model before and after pruning. It can be seen that the indicators such as the number of parameters, the amount of calculation, and the model size have all decreased to varying degrees, achieving the main purpose of pruning and reducing the calculation and storage capabilities required in the training of the network model. In addition, Table 2 lists the segmentation comparison of different evaluation indicators of the U-Net network model before and after pruning in the DRIVE test set. The segmentation comparison diagrams of the DRIVE 4-fold amplified dataset and the 8-fold amplified dataset are respectively as Figure 5 (a) and (b) in

[0052] Table 2 Segmentation results of the U-Net network model before and after pruning on the DRIVE dataset

[0053]

[0054]

Claims

1. A pruning method for U-Net network model based on genetic algorithm, characterized in that, it includes the following steps: Step 1: Horizontally flip, vertically flip and rotate the training set images at multiple angles; Step 2: Perform preprocessing operations on the color images; Step 3: Construct a new loss function for the U-Net network model to make it more suitable for the training of retinal image datasets with pixel imbalance problems; The expression of the new loss function is: Loss=L dice +λL r (1) L r is the cross-entropy function, and its expression is as follows: where TP and TN are the numbers of true positive and true negative pixels respectively; N p and N n are the numbers of target pixels and non-target pixels respectively; y is the label value, y = 1 for the segmentation target, y = 0 for the background; p is the predicted probability value of the pixel; L dice is the Dice coefficient expression, as follows: where N is the number value of pixels; p(k,i) ∈ [0,1], q(k,i) ∈ [0,1] are the predicted probability and true label of pixel point k class respectively, and λ is a coefficient; Step 4: Take the number of four symmetric convolutional layers of the U-Net network model as the optimization target, that is, the optimization result of the genetic algorithm is a combination containing four compression rate values; Use the genetic optimization algorithm to optimize this combination, and finally obtain an optimal compression rate combination. After performing corresponding modification operations on the number of each convolutional layer, a pruned U-Net network model is obtained; Step 5: Incorporate the loss function of the U-Net network into the fitness function of the genetic algorithm to make the two have an inverse relationship, and expand the gap between the loss functions through the fitness function; Step 6: Through continuous iteration of the genetic algorithm, obtain the required optimal compression rate combination.

2. A pruning method for U-Net network model based on genetic algorithm according to claim 1, characterized in that: Horizontally flipping and vertically flipping the training set images amplifies the data volume by 4 times; Rotating the training set images at multiple angles amplifies the data volume by 8 times.

3. A pruning method for U-Net network model based on genetic algorithm according to claim 1, characterized in that: In step 4, the compression rate combination is as follows: α＝[α 1 ，α 2 ，α 3 ，α 4 ] (4) where α 1 , α 2 , α 3 , α 4 ranges from [0.1, 0.2, …, 0.8]; α is the input parameter of the genetic algorithm.

4. A pruning method for U-Net network model based on genetic algorithm according to claim 1, characterized in that: The U-Net network structure is as follows: In the encoding structure of the network, four encoding blocks are used, and each encoding block includes two convolutional layers, one batch normalization layer, LReLU activation layer and max pooling layer; In the decoding structure of the network, corresponding to the encoding structure, four decoding blocks are used, and each decoding block includes two convolutional layers, one upsampling layer, BN layer and one skip connection layer.

5. A pruning method for U-Net network model based on genetic algorithm according to claim 1, characterized in that: In step 5, the expression of the fitness function in the genetic algorithm is: Fitness=e 20*(0.3-Loss) (5) where Loss is the loss function of the U-Net network model.

6. A pruning method for U-Net network model based on genetic algorithm according to claim 1, characterized in that: The pseudo-code for the combination of genetic algorithm and U-Net network model is as follows:

Images