1. A medical image segmentation model compression method is characterized by comprising the following steps:

s1, collecting data in the medical image database;

s2, preprocessing data;

s3, aiming at a medical image segmentation basic model, constructing a search space according to the number of convolution kernels used at each position in the model, and aiming at a coding-decoding structure of a segmentation network, searching a sub-network with small calculation amount and high segmentation precision in the search space by using a symmetric neural network search, wherein the coding-decoding structure is symmetric;

s4, when traversing the whole search space, using a weight sharing method to reduce the calculation cost and the training resource;

and S5, in the training process of the network, a knowledge distillation method is used, the basic model is used as a teacher mode, the compression sub-network is used as a student model, and the knowledge transfer between the basic model and the student model is realized.

2. The medical image segmentation model compression method as claimed in claim 1, wherein the data preprocessing in step S2 includes motion correction, spatial normalization, gray normalization, edge removal, size clipping and center clipping.

3. The medical image segmentation model compression method according to claim 1, wherein the symmetrical neural network search in step S3 includes the following specific steps:

s301, in the split network, a search space is constructed according to the selection of the number of convolution layer channels in the encoding process, and the arrangement of the number of convolution channels in each layer in the search space is { c }₁,c₂,...,c_KK represents the kth network layer to be pruned;

s302, the channel configuration form of the optimal sub-network obtained by network search is as follows:wherein, F_tIs a computational constraint;

s303, generating a corresponding number of convolution channels in the decoding process by utilizing the relation between the convolution access number in the encoding process and the convolution access number in the decoding process, and finally obtaining a subnetwork with a symmetric encoding-decoding structure.

4. The medical image segmentation model compression method according to claim 3, wherein the weight sharing method in step S4 includes the following steps:

s401, based on basic modelA channel, given a network structure configuration

S402, extracting the first weight from the corresponding weight in the one-time-use networkThe channels are used as weights of the structural sub-network;

s403, during each forward propagation in the training process, randomly selecting a sub-network with a certain channel number configuration, calculating the output and gradient of the sub-network, updating the extracted weight according to the learning target, and simultaneously freezing the weights of other parts without participating in the forward propagation;

s404, after the network structure training is finished, traversing the whole search space for testing, and finding a model with the best comprehensive performance;

s405, fine adjustment is carried out by the sub-network to obtain the sub-network with the optimal structure.

5. The compression method for medical image segmentation models according to claim 1, characterized in that the specific steps of knowledge distillation in step S5 are as follows:

s501, adopting a basic model as a teacher model and a compression sub-model as a student model;

s502, training a student model by using a final goal of knowledge distillation;

s503, transferring the intermediate representation of the segmentation model from the teacher model to the student model:

L＝L_seg+λ_distillL_distillwherein L is_segIs a medical image segmentation error, L_distillThe method is the distillation error when the teacher model and the student model transfer knowledge, and the super parameter lambda is used for controlling the importance of the distillation target.

Background

The task of medical image segmentation has been a research hotspot in the fields of computer vision and nature. With the rapid development and application of Convolutional Neural Networks (CNNs), more and more Deep Learning (DL) based medical segmentation models are proposed and have achieved good results in many disease segmentation tasks. On one hand, the number of layers of the neural network is more and more; in another aspect, the development of medical devices also provides higher resolution data. This makes the medical image segmentation task work better and better, but the volume of the model is larger and larger, which does not facilitate the application and deployment of the model in a hardware environment. Therefore, the volume and computational cost of the medical image segmentation model is to be further optimized by model compression.

A good network structure is the key to model compression. At present, the number of convolution kernels and the number of channels of the medical segmentation model are fixed, and a large amount of redundancy is contained in the convolution kernels and the channels. In order to remove the redundancy of the number of convolution channels, it is an alternative to directly and manually choose to reduce the number of convolution kernels, but this can greatly reduce the model performance. While the use of Neural network Search (NAS) can find the structural subnetwork with the best overall performance within the Search space provided by the basic network structure.

At present, in a medical image segmentation task, a symmetric coding-decoding structure is most commonly used, and in order to avoid losing the symmetry in a pruning process, a feature map of a decoding process and a feature map of a coding process are in one-to-one correspondence in scale. Loss of network structure symmetry significantly reduces the segmentation effect compared to maintaining such symmetry during model compression.

Meanwhile, because the structural subnetworks of the search space are all independent deep learning segmentation models, and some local information is repeatedly used in the construction process of each model, a weight sharing strategy can be used, so that the structural subnetworks share some weights, and the time cost and the training resources required by training are reduced. The structural sub-networks obtained by neural network search need to be trained for the medical image segmentation task. However, training a subnetwork ab initio is a difficult result to achieve comparable to a base network, since the size of the subnetwork is smaller than the base network.

Disclosure of Invention

Aiming at the technical problems, the invention provides a medical image segmentation model compression method based on neural network search and knowledge distillation, and a symmetric-NAS coding-decoding structure is introduced to ensure that all sub-networks in a search space have the symmetry. Meanwhile, the sub-network model is similar to the original basic model, the basic model is used as a teacher model, the compression sub-model is used as a student model, the middle representation of the medical image segmentation model is transferred from the teacher model to the student model by using a knowledge-distillation method, the segmentation effect is guaranteed, and meanwhile the calculation cost of the medical image segmentation model is reduced.

In order to achieve the above purpose, the invention provides the following technical scheme:

the invention provides a medical image segmentation model compression method, which comprises the following steps:

s1, collecting data in the medical image database;

s2, preprocessing data;

s3, aiming at a medical image segmentation basic model, constructing a search space according to the number of convolution kernels used at each position in the model, and aiming at a coding-decoding structure of a segmentation network, searching a sub-network with small calculation amount and high segmentation precision in the search space by using a symmetric neural network search, wherein the coding-decoding structure is symmetric;

s4, when traversing the whole search space, using a weight sharing method to reduce the calculation cost and training resources;

and S5, in the training process of the network, a knowledge distillation method is used, the basic model is used as a teacher model, the compression sub-network is used as a student model, and the knowledge transfer between the basic model and the student model is realized.

In the medical image segmentation model compression method, the data preprocessing in step S2 includes motion correction, spatial normalization, gray level normalization, edge removal, size clipping and center clipping.

In the compression method of the medical image segmentation model, the specific steps of the symmetric neural network search in step S3 are as follows:

s301, constructing a search space according to the selection of the number of the convolution layer channels in the coding process in the split network, wherein the arrangement of the number of the convolution channels in each layer in the search space is { c }₁,c₂,...,c_KWhere K denotes the Kth network layer to be pruned, c_iRepresenting the convolution channel number of the ith network layer;

s302, the channel configuration form of the optimal sub-network obtained by network search is as follows:wherein, F_tIs a computational constraint;

s303, generating a corresponding number of convolution channels in the decoding process by utilizing the relation between the convolution access number in the encoding process and the convolution access number in the decoding process, and finally obtaining a subnetwork with a symmetric encoding-decoding structure.

In the compression method of the medical image segmentation model, the specific steps of the weight sharing method in step S4 are as follows:

s401, based on basic modelA channel, given a network structure configuration

S402, extracting the first weight from the corresponding weight in the once-for-all (OFA)The channels are used as weights of the structural sub-network;

s403, during each forward propagation in the training process, randomly selecting a sub-network with a certain channel number configuration, calculating the output and gradient of the sub-network, updating the extracted weight according to the learning target, and simultaneously freezing the weights of other parts without participating in the forward propagation;

s404, after the network structure training is finished, traversing the whole search space for testing, and finding a model with the best comprehensive performance;

s405, fine adjustment is carried out by the sub-network to obtain the sub-network with the optimal structure.

In the compression method of the medical image segmentation model, the specific steps of knowledge distillation in step S5 are as follows:

s501, adopting a basic model as a teacher model and a compression sub-model as a student model;

s502, training a student model by using a final goal of knowledge distillation;

s503, transferring the intermediate representation of the segmentation model from the teacher model to the student model: l ═ L_seg+λ_distillL_distillWherein L is_segIs a medical image segmentation error, L_distillThe method is the distillation error when the teacher model and the student model transfer knowledge, and the super parameter lambda is used for controlling the importance of the distillation target.

Compared with the prior art, the invention has the beneficial effects that:

according to the medical image segmentation model compression method provided by the invention, aiming at a medical image segmentation basic model, a search space is constructed according to the number of convolution kernels used at each position in the model. Aiming at the coding-decoding structure of the segmentation network, a symmetric neural network is used for searching and finding a sub-network with small calculation amount and high segmentation precision in the search space, and the coding-decoding structures are symmetric so as to ensure the segmentation performance. Wherein weight sharing strategies are used to reduce computational cost and training resources while traversing the entire search space. And finally, a knowledge distillation method is used in the training process of the network, the basic model is used as a teacher mode, the compression sub-network is used as a student model, and the knowledge transfer between the basic model and the student model is realized. According to the invention, through neural network search and knowledge distillation, the calculation cost for constructing the network is greatly reduced on the premise of ensuring the segmentation effect of the medical image segmentation model, the model structure is optimized, and the method can be applied to various medical image segmentation models.

Drawings

In order to more clearly illustrate the embodiments of the present application or technical solutions in the prior art, the drawings needed to be used in the embodiments will be briefly described below, and it is obvious that the drawings in the following description are only some embodiments described in the present invention, and it is obvious for those skilled in the art to obtain other drawings based on these drawings.

Fig. 1 is a schematic diagram of model compression using Res-Unet network as a basic model according to an embodiment of the present invention.

Fig. 2 is a diagram illustrating an optimized result of brain tumor lesion segmentation according to an embodiment of the present invention.

Detailed Description

For a better understanding of the present solution, the method of the present invention is described in detail below with reference to the accompanying drawings.

The invention provides a medical image segmentation model compression method, which comprises the following steps:

step S1, data in the medical image database is collected, in this embodiment, taking magnetic resonance images of brain tumor patients as an example, the four modalities are mainly T1, T1c, T2 and FLAIR.

And step S2, performing data preprocessing, including motion correction, spatial standardization, gray scale normalization, dehulling and neck removal, and size cutting. Center cropping was then performed on each 3D MRI tested, preserving the entire brain region, and removing the border black regions.

And step S3, using Res-Unet as basic skeleton of network, and using separable convolution as convolution layer. Res-uet is a residual U-type network. The Unet is a classical network for medical image segmentation, and in the embodiment, a Residual module is introduced into a network structure to form a Res-Unet network.

First, in the network, a search space is constructed according to the selection of the number of convolutional layer channels in the encoding process, in this embodiment, the number of convolutional channels is selected to be a multiple of 8, and possible channel configurations are { c }₁,c₂,...,c_KWhere K denotes the number of network layer for pruning. The channel configuration of the optimal sub-network is as follows:

wherein, F_tAre computational constraints.

Then, the corresponding number of convolution channels is generated in the decoding process by utilizing the relation between the number of convolution channels in the encoding process and the number of convolution channels in the decoding process. Finally, a subnetwork which is small in calculation amount, high in segmentation precision and symmetrical in coding-decoding structure is obtained.

And step S4, traversing the whole search space by using a weight sharing method in the process of training the network.

Specifically, each time forward propagation occurs, a subnet is randomly selected for activation, while other weights are frozen. Each subnetwork has an equal chance of being selected and trained.

Assume that the basic model hasA number of channels configured for a given number of channelsBy extracting the first one from the corresponding weight tensor in an once-for-all (OFA) networkThe channel serves as a weight for the sub-network.

During each forward propagation of the training process, a sub-network with a certain channel number configuration is randomly selected, the output and the gradient of the sub-network are calculated, the extracted weight is updated according to a learning target, and other weights are frozen and do not participate in the forward propagation. After the network structure training is completed, the model with the best comprehensive performance can be found only by traversing the whole search space for testing. And finally, fine tuning is carried out by utilizing the sub-network, and the purpose of searching the neural network is realized.

Step S5, using knowledge distillation to improve compression sub-model performance. And (3) adopting the basic model as a teacher model and the compression sub-model as a student model, and transferring the intermediate representation of the segmentation model from the teacher model to the student model. The goals of knowledge distillation are:

wherein S is_t(x) And S'_t(x) Is the intermediate characteristic of the T-th layer selected from the student model and the teacher model, T is the number of layers, f_tIs a 1x1 convolutional layer for mapping the features of the student model to the same number of channels in the corresponding features of the teacher model. The final goals are:

L＝L_seg+λ_distillL_distill

among these, the super-parameter λ is used to control the importance of distillation targets.

The model compression diagram using Res-Unet network as basic model is shown in FIG. 1. Fig. 1 shows two Res-Unet networks, which can be seen as a symmetric codec structure network. The Res-Unet network on the upper side is a basic model and is used as a teacher model; the Res-Unet network below is a structural sub-network obtained by weight sharing and neural network searching, and serves as a student model. The intermediate representation of the medical image segmentation model is migrated from the teacher model to the student model with the goal of minimizing distillation loss through knowledge distillation.

The optimization results for the example of brain tumor lesion segmentation are shown in fig. 2. The figure contains 4 cases, and each case comprises a nuclear magnetic image of the brain tumor, a brain tumor focus label, a Res-Unet basic network segmentation result and a sub-network segmentation result after model compression from left to right. As can be seen from the figure, compared with the focus label, the segmentation performance of the basic network and the network after model compression is approximately consistent, but the network parameter after model compression is far smaller than that of the basic network, which shows that the invention can greatly reduce the model amount on the premise of ensuring the segmentation performance of the medical image.

The invention designs a medical image segmentation model compression method based on neural network search and knowledge distillation. Firstly, a sub-network with the best performance and a symmetrical structure of a medical image segmentation network is searched and found based on a symmetrical neural network, then, weight sharing strategies are used for optimizing and calculating cost when the network is trained, and meanwhile, the intermediate representation of a teacher model is migrated to a student model based on knowledge distillation, so that the calculation performance is further optimized. The method integrates different strategies, can greatly reduce the calculated amount of the model and optimize the model structure on the premise of ensuring the segmentation effect of the medical image, and facilitates the landing and the deployment of the subsequent model in the actual application.

The above examples are only intended to illustrate the technical solution of the present invention, but not to limit it; although the present invention has been described in detail with reference to the foregoing embodiments, it will be understood by those of ordinary skill in the art that: it is to be understood that modifications may be made to the technical solutions described in the foregoing embodiments, or equivalents may be substituted for some of the technical features thereof, but these modifications or substitutions do not substantially depart from the spirit and scope of the technical solutions of the embodiments of the present invention.