Deep learning model based on attention mechanism

Abstract

The application discloses a deep learning model based on an attention mechanism, comprising a data set, the data set is proportionally divided into a training data set and a test data set after scaling; the deep learning model is composed of an encoder responsible for extracting feature information and a decoder responsible for recovering spatial information.The application combines convolution operation with self-attention mechanism, extracts local low-level features and global context features at the same time, introduces channel and spatial attention modules on the skip connection between the encoder and the decoder to suppress irrelevant information, maximally utilizes useful information, thereby improves segmentation precision, and realizes more accurate brain hematoma segmentation.

Description

TECHNICAL FIELD

[0001] The application relates to the technical field of image segmentation, in particular to a deep learning model based on an attention mechanism. BACKGROUND

[0002] With the wide application of deep learning technology in the field of medical image processing in recent years, a semantic segmentation model based on a convolutional neural network has excellent performance in segmenting a target lesion region. The feature extraction capability and the spatial information recovery capability of the model play an important role in segmentation accuracy and directly affect the final prediction result.

[0003] In a brain hematoma segmentation task, the hematoma in a brain CT image is a high-density area, but due to the complexity of the brain structure, the diversity of the hematoma shape and position, it is extremely difficult to accurately and reliably segment the hematoma. The convolution operation has the limitations of only being able to perform linear operation and extract local features, and the segmentation task often has over-segmentation and under-segmentation problems, especially irregularly shaped hematoma and hematoma close to the skull. Therefore, it is difficult to accurately segment the brain hematoma by using the existing semantic segmentation model. SUMMARY

[0004] The application aims to solve the problems in the prior art and provides a deep learning model based on an attention mechanism, which uses a self-attention mechanism, a channel attention mechanism and a spatial attention mechanism to improve segmentation accuracy.

[0005] To achieve the above purpose, the application adopts the following technical scheme:

[0006] The deep learning model based on the attention mechanism comprises:

[0007] A data set, which is divided into a training data set and a test data set in proportion after scaling;

[0008] The deep learning model comprises an encoder responsible for extracting feature information and a decoder responsible for recovering spatial information.

[0009] Preferably, the deep learning model is specifically designed as follows:

[0010] (2.1) The encoder performs four times of down-sampling max-pooling operation and comprises four layers of residual networks and a self-attention network, the residual networks are used to avoid the gradient disappearance problem and improve the propagation rate of features, and the self-attention network is located at the bottom of the encoder and extracts the feature correlation characteristics of the whole image through high-order operation to extract the hematoma features from a global perspective;

[0011] (2.2) Decoder performs four times up-sampling bilinear interpolation operation, which consists of four convolutional layers and the last classification layer, each convolutional layer includes two groups of 3*3 convolution, Batch Normalization and nonlinear activation function ReLU, and the classification layer consists of 3*3 convolution and sigmoid activation function;

[0012] (2.3) There is a skip connection between the encoder and the decoder, the encoder features pass through the channel attention module and the spatial attention module to learn the weights on each channel and feature area, highlight the useful feature information and suppress irrelevant information, and then perform channel splicing with the up-sampling output of the previous layer of the decoder to improve the spatial information recovery efficiency.

[0013] Preferably, the self-attention network is composed of multiple attention heads located at the bottom of the encoder. Its main role is to obtain the receptive field of the entire input image by establishing a connection between each pixel in the high-level feature map. Therefore, the classification decision of a specific pixel of the input image when segmenting the brain hematoma region will be affected by any other pixel. The self-attention calculation formula is as follows:

[0014]

[0015] Preferably, the channel attention module increases the feature map X generated by the input CT image through the convolutional layer from single channel to multiple channels. The information expressed by the feature map of each channel is different, and the effective feature information may only appear in a specific channel. The role of the channel attention module is to learn the weights by using the relationship between each channel, and then multiply the corresponding channel.

[0016] Preferably, the calculation formula of the channel attention coefficient is as follows:

[0017] Attention C (X)=σ(MLP(Avgpool(X))+MLP(Maxpool(X)))

[0018] In the formula, σ represents the sigmoid activation function, MLP represents the multi-layer perceptron, and Avgpool and Maxpool represent the global average pooling and the global maximum pooling, respectively.

[0019] Preferably, the spatial attention module focuses on the information most meaningful to the current segmentation task, and the calculation formula of the spatial attention coefficient is as follows:

[0020] Attention S (X)=σ(f 7×7 ([Avgpool(X);Maxpool(X)]))

[0021] wherein sigma represents a sigmoid activation function, f 7×7 represents a 7*7 convolution operation, Avgpool and Maxpool represent global average pooling and global maximum pooling respectively.

[0022] Preferably, the training data set and the test data set are in a ratio of 9:1.

[0023] The advantage of the present application is that the deep learning model based on the attention mechanism provided by the present application extracts local low-level features and global context features by combining convolution operation and self-attention mechanism. At the same time, the channel and spatial attention modules are introduced on the skip connection between the encoder and the decoder to suppress irrelevant information and maximize the use of useful information, thereby improving the segmentation accuracy and realizing more accurate brain hematoma segmentation. BRIEF DESCRIPTION OF DRAWINGS

[0024] Figure 1 is a deep learning model structure diagram of the attention mechanism of the present application;

[0025] Figure 2 is a self-attention network structure diagram of the present application;

[0026] Figure 3 is a channel attention module structure diagram of the present application;

[0027] Figure 4 is a spatial attention module structure diagram of the present application. DETAILED DESCRIPTION

[0028] In order to make the purpose, technical scheme and advantages of the present application more clear, the present application is further described in detail below in combination with the drawings and examples. It should be understood that the specific examples described herein are only used to explain the present application and do not limit the present application.

[0029] As shown in Figures 1-4 , the deep learning model based on the attention mechanism provided by the present application comprises:

[0030] a data set, the data set is scaled and divided into a training data set and a test data set in proportion;

[0031] a deep learning model, which is composed of an encoder responsible for extracting feature information and a decoder responsible for restoring spatial information.

[0032] The specific steps are as follows:

[0033] (1) Collect brain CT or brain MR images as a data set, scale and divide into a training data set and a test data set in a ratio of 9:1;

[0034] (2) a deep learning model for constructing an attention mechanism, composed of an encoder responsible for extracting feature information and a decoder responsible for recovering spatial information;

[0035] (3) iteratively training the deep learning model with the training set as input until the loss converges;

[0036] (4) inputting the test set into the model trained in step (3) and outputting the final prediction result.

[0037] The deep learning model of step (2) is divided into two parts, encoder and decoder, and the specific design is as follows:

[0038] (2.1) The encoder performs four times of down-sampling max-pooling operation, which is composed of four layers of residual network and self-attention network. The residual network is used to avoid the problem of gradient disappearance and improve the propagation rate of features. The self-attention network is located at the bottom of the encoder, which extracts the feature correlation characteristics of the whole image through high-order operation, and extracts the hematoma features from a global perspective;

[0039] (2.2) The decoder performs four times of up-sampling bilinear interpolation operation, which is composed of four layers of convolution layer and the last classification layer. Each convolution layer includes two groups of 3*3 convolution, batch normalization and nonlinear activation function ReLU. The classification layer is composed of 3*3 convolution and sigmoid activation function;

[0040] (2.3) There is a skip connection between the encoder and the decoder. The encoder features pass through the channel attention module and the spatial attention module to learn the weights on each channel and feature area. Through the weights, useful feature information is highlighted and irrelevant information is suppressed. Then, the output of the up-sampling of the previous layer of the decoder is channel spliced to improve the spatial information recovery efficiency.

[0041] The input of the deep learning model is a single-channel original gray image of 512*512, and the local features are extracted through the convolution layer in the residual network, while the number of feature map channels is doubled. The downsampling is a max-pooling operation, and the size of the feature map is reduced to half of the original size after each downsampling. The self-attention network at the end of the encoder expands the receptive field to the whole image, and further utilizes the global context information. The convolution layer of each layer of the decoder is connected to the corresponding encoder layer through the channel and spatial attention modules. The output from the channel and spatial attention modules is channel spliced with the up-sampling output of the previous layer of the decoder. The decoder layer by layer increases the size of the feature map by one time through the bilinear interpolation up-sampling and the convolution layer, and reduces the number of channels by half to restore the spatial information of the image. The final result of the up-sampling is a single-channel feature map of the original input size, and then the classification layer composed of 3*3 convolution and sigmoid activation function is used to output the segmentation result. All the convolution layers are composed of 2 (Conv3*3+BN+ReLU), wherein Conv3*3 is a 3*3 convolution, BN is batch normalization, and ReLU is a nonlinear activation function.

[0042] Self-attention network:

[0043] The self-attention network is composed of multiple attention heads located at the bottom of the encoder. Its main function is to obtain the receptive field of the whole input image by establishing a connection between each pixel in the high-level feature map. Therefore, when segmenting the brain hematoma region, the classification decision of a specific pixel of the input image will be affected by any other pixel. The self-attention calculation formula is as follows:

[0044]

[0045] As shown in Figure 2 , the input of the self-attention network is a feature map X with a height of H, a width of W and a channel number of C. In order to consider the global context information, the position encoding information is added to X. The position encoding is of great significance for brain hematoma segmentation, because different brain tissues are located at different fixed positions in the CT image. After position encoding, the absolute position and relative position information between brain tissues can be captured. Then X is reshaped into three one-dimensional vectors, which are query matrix Q, key matrix K and value matrix V. A is the attention coefficient matrix, and its actual meaning is the correlation between a given element in Q and all elements in K. The weighted average of the elements in the value matrix is calculated according to the self-attention formula to obtain the final attention output, and the output feature map Y is reshaped. In the brain hematoma segmentation task of the present application, the dimensions of Q, K and V are the same, which are the results of one-dimensional vectors after different embedding. The embedding matrix is represented as Wq, Wk and Wv.

[0046] Channel attention module:

[0047] AsFigure 3 As shown, the channel attention module increases the number of channels in the feature map X generated from the input CT image through convolutional layers from a single channel to multiple channels. Each channel's feature map expresses different information, and effective feature information may only appear in a specific channel. The function of the channel attention module is to learn weights using the relationships between each channel, and then multiply these weights by the corresponding channel.

[0048] The formula for calculating the channel attention coefficient is as follows:

[0049] Attention C (X)=σ(MLP(Avgpool(X))+MLP(Maxpool(X)))

[0050] In the formula, σ represents the sigmoid activation function, MLP represents the multilayer perceptron, and Avgpool and Maxpool represent global average pooling and global max pooling, respectively.

[0051] Spatial attention module:

[0052] like Figure 4 As shown, the spatial attention module focuses on the information most meaningful for the current segmentation task. In CT images, brain hematoma regions exhibit blurred boundaries and low contrast. Therefore, spatial attention is utilized on skip connections to improve the efficiency of spatial information aggregation. The role of the spatial attention module is to model the priority relationships of spatial locations. For example... Figure 4 As shown, to effectively learn spatial attention weights, the spatial attention module uses global average pooling and global max pooling operations to reduce the dimensionality of the input feature map X, generating two feature maps for each spatial location. These two maps are then concatenated channel-wise and further reduced to a single channel via a 7x7 convolution. Finally, the sigmoid function is used to activate the weights, which are then multiplied by the input to obtain the final spatial attention output. The formula for calculating the spatial attention coefficients is as follows:

[0053] Attention S (X)=σ(f 7×7 ([Avgpool(X);Maxpool(X)]))

[0054] In the formula, σ represents the sigmoid activation function, f 7×7 This represents a 7x7 convolution operation. Avgpool and Maxpool represent global average pooling and global max pooling, respectively.

[0055] While embodiments of the application have been shown and described, it is to be understood that the embodiments described are merely exemplary of the principles and application of the present application. Numerous modifications and adaptions can be effected without departing from the spirit and scope of the present application, which is not limited to the exact construction and arrangement described. It is intended, therefore, to cover all modifications and adaptions that fall within the scope of the claims and their equivalents.

Claims

1. A method for constructing a deep learning model based on an attention mechanism, characterized in that, Comprise: a data set, the data set is divided into training data set and test data set after scaling by proportion; a deep learning model, which is composed of an encoder responsible for extracting feature information and a decoder responsible for recovering spatial information; the deep learning model is specifically designed as follows: (2.1) the encoder performs four times of downsampling max-pooling operation, which is composed of four layers of residual network and self-attention network, the residual network is used to avoid gradient disappearance problem and improve the propagation rate of features, the self-attention network is located at the bottom of the encoder, which extracts the feature correlation characteristics of the whole image through high-order operation, and extracts the hematoma features from the global perspective; (2.2) the decoder performs four times of upsampling bilinear interpolation operation, which is composed of four layers of convolution layer and the last classification layer, each convolution layer includes two groups of 3*3 convolution, batch normalization and nonlinear activation function ReLU, and the classification layer is composed of 3*3 convolution and sigmoid activation function; (2.3) there is a skip connection between the encoder and the decoder, the encoder features pass through the channel attention module and the spatial attention module to learn the weight on each channel and feature area, highlight the useful feature information and suppress irrelevant information, then the channel splicing is carried out with the up-sampling output of the previous layer of the decoder to improve the spatial information recovery efficiency. 2.The method of claim 1, wherein: The self-attention network is composed of multiple attention heads at the bottom of the encoder, which mainly establishes the connection between each pixel in the high-level feature map to obtain the receptive field of the whole input image, so when segmenting the brain hematoma area, the classification decision of a certain pixel of the input image will be affected by any other pixel, the self-attention calculation formula is as follows: 。 3.The method of claim 1, wherein: The channel attention module increases the feature map X generated by the input CT image through the convolution layer from single channel to multiple channels, the information expressed by the feature map of each channel is different, and the effective feature information may only appear in a specific channel, the role of the channel attention module is to learn the weight by using the relationship between each channel, and then multiply the corresponding channel. 4.The method of claim 3, wherein: The calculation formula of channel attention coefficient is as follows: wherein denotes a sigmoid activation function, MLP denotes a multi-layer perceptron, Avgpool and Maxpool denote global average pooling and global max pooling, respectively. 5.The method of claim 1 or 2 or 3 or 4, wherein: The spatial attention module focuses on the information most meaningful to the current segmentation task, and the calculation formula of spatial attention coefficient is as follows: wherein denotes a sigmoid activation function, denotes a 7*7 convolution operation, Avgpool and Maxpool denote global average pooling and global max pooling, respectively. 6.The method of claim 1 to 4, wherein: The proportion of training data set and test data set is 9:1.

Images