Medical image segmentation model based on lightweight attention module and training method of model

Abstract

This invention discloses a medical image segmentation model and its training method based on a lightweight attention module. The segmentation model is constructed based on a UNet network and includes 4-6 upsampling and downsampling stages. Each upsampling and downsampling stage includes a convolutional sampling module. The downsampling layer uses pooling operations, and the upsampling layer uses interpolation operations. Each convolutional layer is followed by a batch normalization layer. A local spatial attention module formed by a combination of large-kernel dilated convolutions and a channel attention module that interacts between channels are added in series in the downsampling layer. Convolutional gated attention modules are added to skip connections and the upsampling layer. A global attention module is set at the network bottleneck. The convolutional attention module in this invention enhances the model's multi-scale and multi-directional learning and understanding of images at the local-global level, resulting in higher segmentation accuracy, stronger robustness, simpler training methods, and faster running speed.

Description

Technical Field

[0001] This invention relates to the field of image processing technology, specifically to a medical image segmentation model based on a lightweight attention module and a training method for the model. Background Technology

[0002] Medical image segmentation is a crucial task in medical image analysis, aiming to separate lesion information from the background to support doctors in qualitative and quantitative assessment and diagnosis of diseases. In recent years, the use of medical image segmentation in point-of-care medical devices has been increasing. For example, point-of-care ultrasound (POCUS) devices have proven useful for rapidly examining pleural irregularities and automatically calculating bladder volume. Images from mobile phone cameras are also used to detect and diagnose skin conditions. These tasks are primarily solved using the UNet deep learning approach. This is because UNet innovatively introduces skip connections on top of fully convolutional networks, fully utilizing multi-stage feature information in downsampling while maintaining network simplicity, resulting in powerful segmentation performance. However, convolutional operations typically exhibit limitations when acquiring relationships between long-range information.

[0003] Unlike convolution, ViT (Vision Transformer) better integrates contextual information by labeling images as sequences and utilizing self-attention blocks to facilitate global communication. However, the self-attention module has O(d^2) parameters per head, where d is the dimension of the input, which is particularly demanding on the size of the input for global computation. Therefore, this type of network sacrifices model structural simplicity for improved segmentation accuracy, making it unsuitable for clinical applications in constrained medical settings.

[0004] As the trend of exploring how to achieve high performance with small models gradually rises, designing an efficient network with attention performance that has less computational overhead, fewer parameters, and faster inference time, while maintaining good performance, is crucial to meeting the real-time application needs of the medical field. Summary of the Invention

[0005] To address the shortcomings of the existing technologies, the technical problem to be solved by this invention is: how to provide a medical image segmentation model based on a lightweight attention module that has high segmentation accuracy, relatively few operating parameters and computational load, and faster operating speed.

[0006] To solve the above-mentioned technical problems, the present invention adopts the following technical solution:

[0007] A medical image segmentation model based on a lightweight attention network is proposed. The segmentation model is constructed using a UNet network and includes 4-6 upsampling / downsampling stages. Each upsampling / downsampling stage includes a convolutional sampling module. The downsampling layer in each stage employs pooling, while the upsampling layer employs interpolation. Each convolutional layer is followed by a batch normalization layer. A local spatial attention module formed by a combination of large-kernel dilated convolutions and a channel attention module for interaction between channels are added in series within the downsampling layer. The local spatial attention module generates attention coefficients for different features in the local convolutional space, while the channel attention module compresses information along the channel direction for feature learning in the channel dimension. Convolutional gated attention modules are added to the skip connections and upsampling layers to filter irrelevant information and ensure the effectiveness of feature upsampling. A global attention module is placed at the network bottleneck to calculate the correlation between each spatial location and all other spatial locations globally, realizing long-range dependencies between pixel locations.

[0008] As an optimization, the segmentation model includes five upsampling and downsampling stages. The number of upsampling and downsampling stages is adjusted to meet the requirement of lightweight design without affecting performance.

[0009] As an optimization, the local spatial attention module uses 1×1 convolution and the GELU activation function to activate the original feature parameters. Then, it uses three large-kernel dilated convolutions of different sizes (9×9, 13×13, and 21×21) to extract image features at multiple scales, which are then fused to generate an attention weight map. Finally, a residual structure is added to connect the input to the dot product of the attention map. The specific steps are as follows:

[0010] x1,x2,x3=GELU(Conv1(x)) (1)

[0011] Att = Conv1(D9(x1) + D 13 (x2)+D 21 (x3)) (2)

[0012] Out = x + Conv1(Att.x) (3)

[0013] Where x represents the input feature information, D represents large kernel dilated convolution, Conv represents ordinary convolution, Att represents the attention weight, and GELU represents the activation function;

[0014] The channel attention module uses average pooling to generate channel features, replacing depthwise convolutions with fully connected layers. Leveraging their highly non-linear global effects, it adaptively generates channel attention maps, as shown in the following formula:

[0015] Att=σ(Linear(RELU(Linear(AvePool(x))))) (4)

[0016] out = x + x.Att (5)

[0017] Where σ represents the sigmoid activation function, Linear represents the fully connected structure, ReLU represents the activation function, and AvePool is the average pooling operation.

[0018] As an optimization, the gating attention module extracts information from the coarse scale as approximate features of the input tensor. This information is then multiplied by the detailed features extracted from the finer scale after being activated by different activation functions. This process performs a secondary feature selection, giving the final weight map a detailed classification capability and accurately identifying the usefulness of information. The formula for the attention module is shown below:

[0019] F1=σ(DW3(x acted (6)

[0020] F2 = GELU(DW5(x) acted (7)

[0021] Att = GELU(DW3(F1.F2)) (8)

[0022] Out = Conv1(x acted )+Att (9)

[0023] Where (F1, F2) is the feature weight matrix, σ is the sigmoid activation function, and DW is the depthwise separable convolution.

[0024] As an optimization, the global attention module multiplies the input features by their transpose matrix. Each point in the resulting attention map is represented as the sum of the products of all channels at two spatial locations. Following the steps of matrix multiplication, the correlation between the spatial domain represented by each row and the global information can be calculated. The specific formula is as follows:

[0025]

[0026] z i =.y i +x i (11)

[0027] Where i is the position to be calculated, j is the position to be used for weighting, i.e., the position of the feature map; x is the input, y is the output of the dot product, f is the function for calculating similarity, g is a univariate, and C is the normalization; for the original feature map x, a 1×1 convolution is first used to generate θ(x iThen, the similarity relationship f(x) is obtained through softmax. i ,x j ), and combine it with g(x) j Multiply to obtain y i Perform another 1×1 convolution W Z Finally with x i Adding them together gives the final output z i .

[0028] This invention also discloses a training method for a medical image segmentation model, comprising the following steps:

[0029] (1) Obtain the image set, normalize the content of all images in the image set to a value with an average value of 0 and a standard value of 1, and adjust the resolution to 256×256. Then divide the image set into a training set and a validation set.

[0030] (2) Set the hyperparameters and loss function, initialize the optimizer, and initialize the LA-UNet neural network for image segmentation. The LA-UNet neural network is the medical image segmentation model based on the lightweight attention module described above.

[0031] (3) Input the training set x into the LA-UNet neural network to obtain the segmentation result y, calculate the loss between the predicted result and the true value, and then backpropagate to update the network parameters;

[0032] (4) Validate the LA-UNet neural network after updating the network parameters in step (3) using the validation set. Calculate the Dice coefficient and IOU of the current model on the validation set. If the Dice coefficient is better than the historical maximum value, save the current model parameters and update the historical best Dice value.

[0033]

[0034]

[0035] (5) Training ends when the number of iterations reaches the set value.

[0036] As an optimization, the loss function uses a combination of cross-entropy and dice loss, and its formula is:

[0037]

[0038]

[0039] L Ours =0.5L BCE +L Dice (16)

[0040] Where N is the total number of samples, y i p is the label value. i Here, |X| and |Y| represent the predicted values, respectively.

[0041] Compared to existing technologies, this invention offers the following advantages: The convolutional attention module enhances the model's multi-scale and multi-directional learning and understanding of images at the local-global level; compared to the traditional UNet segmentation network, it achieves higher segmentation accuracy and stronger robustness, adapting to multimodal data; compared to U-shaped networks incorporating Transformers, it does not require a large training set to stack model performance, significantly reducing model parameters and computational complexity, simplifying training methods, increasing running speed, and even surpassing segmentation accuracy to a certain extent. This method has significant direct application value for real-time medical auxiliary diagnosis. Attached Figure Description

[0042] Figure 1 This is a diagram showing the overall structure of the model in this invention;

[0043] Figure 2 This is a structural diagram of the local spatial attention module in this invention;

[0044] Figure 3 This is a structural diagram of the channel attention module in this invention;

[0045] Figure 4 This is a structural diagram of the gating attention module in this invention;

[0046] Figure 5 This is a structural diagram of the global attention module in this invention;

[0047] Figure 6 This is a flowchart of the training-testing process for the model in this invention;

[0048] Figure 7 Lightweight metrics for each model (dot plot);

[0049] Figure 8 This is a diagram showing the segmentation structure of the model in this invention;

[0050] Figure 9 This is an attention heatmap during the image processing process of the model in this invention. Detailed Implementation

[0051] like Figure 1As shown in the figure, the medical image segmentation model based on a lightweight attention module in this specific embodiment is constructed based on the UNet network. The segmentation model includes five upsampling and downsampling stages, with spatial resolution decreasing / increasing (16, 32, 128, 160, 256). To achieve lightweight processing, redundant convolutional structures in each layer of the original network are removed. Each upsampling and downsampling stage includes only one convolutional sampling module. Each stage consists of a 3×3 convolution and a series of accompanying functions. The downsampling layer in the upsampling and downsampling stages uses pooling operations, and the upsampling layer uses interpolation operations. Furthermore, each convolutional layer is followed by a batch normalization layer because batch normalization has been found to achieve better segmentation performance than layer normalization. In optimizing the segmentation model, this invention replaces Tra with a lightweight plugin. The attention head in nsformer focuses on contextual attention. In the downsampling layer, a local spatial attention module formed by combining large-kernel dilated convolutions and a channel attention module for interaction between channels are cascaded. The local spatial attention module generates attention coefficients for different features in the local convolutional space, while the channel attention module compresses information along the channel direction for feature learning in the channel dimension. Convolutional gated attention modules are added to skip connections and upsampling layers to filter irrelevant information and ensure the effectiveness of feature upsampling. A global attention module is placed at the network bottleneck to calculate the correlation between each spatial location and all other spatial locations globally, realizing long-range dependencies between pixel locations.

[0052] The Local Spatial Attention (LSA) module aims to generate attention coefficients for different features within a local convolutional space. Mimicking the classic pyramid structure, the LSA module uses 1×1 convolutions and the GELU activation function to activate existing feature parameters. Then, it extracts image features at multiple scales using three large-kernel dilated convolutions of different sizes (9×9, 13×13, and 21×21), which are then fused to generate an attention weight map. Finally, a residual structure connects the input to the dot product of the attention map, as shown below. Figure 2 As shown, the specific steps are as follows:

[0053] x1,x2,x3=GELU(Conv1(x)) (1)

[0054] Att = Conv1(D9(x1) + D 13 (x2)+D 21 (x3)) (2)

[0055] Out = x + Conv1(Att.x) (3)

[0056] Where x represents the input feature information, D represents large kernel dilated convolution, Conv represents ordinary convolution, Att represents the attention weight, and GELU represents the activation function;

[0057] The channel attention module needs to compress information along the channel direction and then perform feature learning along the channel dimension, such as... Figure 3 As shown, the channel attention module uses average pooling to generate channel features. Max pooling is too biased in covering channel information, losing much of its meaning, which also makes the module more lightweight. This application, influenced by the Multilayer Perceptron (MLP), replaces depthwise convolutions with fully connected layers, utilizing their highly nonlinear global effects to ensure low computational cost while fully mixing features from each channel, adaptively generating channel attention maps. The specific formula is as follows:

[0058] Att=σ(Linear(RELU(Linear(AvePool(x))))) (4)

[0059] out = x + x.Att (5)

[0060] Where σ represents the sigmoid activation function, Linear represents the fully connected structure, ReLU represents the activation function, and AvePool is the average pooling operation.

[0061] Gated attention module such as Figure 4 As shown, this process filters out irrelevant information, ensuring the effectiveness of feature upsampling. To reduce the number of trainable parameters and computational complexity, a linear transformation (1×1 convolution) is applied to the input information, and the transformed features are downsampled using convolutions at different depths to achieve the resolution of the gated signal, similar to a local operation. The information extracted from the coarse scale (5×5) by the gated attention module is used as the approximate features of the input tensor. This information is then multiplied by the detailed features extracted from the finer scale after being activated by different activation functions, achieving a secondary feature filtering process. This allows the final weight map to have detailed classification capabilities, accurately identifying the usefulness of information. The formula for the attention module is shown below:

[0062] F1=σ(DW3(x acted (6)

[0063] F2 = GELU(DW5(x) acted (7)

[0064] Att = GELU(DW3(F1.F2)) (8)

[0065] Out = Conv1(x acted )+att (9)

[0066] Where (F1, F2) is the feature weight matrix, σ is the sigmoid activation function, and DW is the depthwise separable convolution.

[0067] Standard convolution operations have limited ability to represent information within a local receptive field; for example, a 3×3 convolution kernel can only process information within a very small range. For a given image, it is difficult to effectively capture long-range dependencies between two points that are far apart. The Global Spatial Attention Module (GSA) aims to compute the correlation between each spatial location and all other spatial locations globally, realizing long-range dependencies between pixel locations. Its structure is as follows: Figure 5 As shown. In the high-level semantic information at the end of the downsampling layer, the feature parameters of a row often map to a region in space. The global attention module multiplies the input features with their transpose matrix, and each point in the attention map is represented as the sum of the products of all channels of two locations in space. This data is the correlation degree. Following the steps of matrix multiplication, the correlation degree between the spatial domain represented by each row and the global information can be calculated. The specific formula is as follows:

[0068]

[0069] z i =.y i +x i (11)

[0070] Where i is the position to be calculated, j is the position to be used for weighting, i.e., the position of the feature map; x is the input, y is the output of the dot product, f is the function for calculating similarity, g is a univariate (1×1 convolution), and C is the normalization; for the original feature map x, a 1×1 convolution is first used to generate θ(x i Then, the similarity relationship f(x) is obtained through softmax. i ,x j ), and combine it with g(x) j Multiply to obtain y i To achieve the same dimension as the original mapping x, another 1×1 convolution W is performed. Z Finally with x i Adding them together gives the final output z i .

[0071] like Figure 6 As shown, a training method for a medical image segmentation model includes the following steps:

[0072] (1) Obtain a set of breast ultrasound images (BUSI), normalize all images in the set to a mean of 0 and a standard value of 1, adjust the resolution to 256×256, use only benign and malignant images, and select 350 images as the training set, 42 images as the validation set, and 45 images as the test set.

[0073] (2) Set hyperparameters and loss function, initialize the optimizer, and initialize the model in this application. All experiments were implemented on a single NVIDIA GeForce RTX3080 GPU. Each dataset was trained for 300 epochs with a batch size of 4 and an initial learning rate of 0.0001. The optimizer was Adam. The learning rate variation strategy was Cosine AnnealingLR with a decay coefficient of 0.9. A cosine annealing learning rate scheduler was used, with a minimum learning rate of 0.00001. A combination of cross-entropy and dice loss was used as the loss function during training. The formula for the loss function is:

[0074]

[0075]

[0076] L Ours =0.5L BCE +L Dice (16)

[0077] Where N is the total number of samples, y i p is the label value. i The values ​​are predicted, where |X| and |Y| represent the label and prediction, respectively.

[0078] (3) Input the training set x into the model to obtain the segmentation result y, calculate the loss between the predicted result and the true value, and then backpropagate to update the network parameters;

[0079] (4) Validate the model after updating the network parameters in step (3) using the validation set. Calculate the Dice coefficient and IOU of the current model on the validation set. If the Dice coefficient is better than the historical maximum value, save the current model parameters and update the historical best Dice value.

[0080]

[0081]

[0082] (5) Training ends when the number of iterations reaches the set value.

[0083] In this specific implementation, the performance of the medical image segmentation model is tested:

[0084] (a) Preprocessing images of each modal test set;

[0085] (b) Select the trained model parameters and load LA-UNet (the model in this application), and calculate the number of model parameters and computational complexity;

[0086] (c) Input the test image, output the segmentation results, and calculate various segmentation metrics based on the corresponding labels: Dice, IOU, Sensitivity, Precision, and Specificity. The model's performance on the BUSI, ISIC, Kvasir, and LITS datasets is recorded in Tables 1 to 4. Table 1 shows the segmentation performance values ​​of each model on the ISIC dataset, Table 2 shows the segmentation performance values ​​of each model on the BUSI dataset, Table 3 shows the segmentation performance values ​​of each model on the Kvasir dataset, and Table 4 shows the segmentation performance values ​​of each model on the LITS dataset. The tables show that LA-UNet, while having small parameters and fast computation, outperforms many U-shaped networks. To demonstrate the lightweight nature of LA-UNet, a dot plot was drawn for its parameter count, computational complexity, and inference time, as shown below. Figure 7 As shown, the vertical axis represents the segmentation accuracy of the model. It can be seen more intuitively that LA-UNet is always in the upper left corner of the graph, indicating that it has high performance. Figure 8 This is a comparison of the LA-UNet test results with the labels; the two are almost identical, demonstrating its segmentation accuracy. During image processing, we extract attention weights at the attention modules of the downsampling layer, bottleneck layer, and upsampling layer to generate corresponding heatmaps, such as... Figure 9 As shown, the network's attention gradually focuses on the segmentation target as the model's inference evolves, which explains the model's advantage in segmentation accuracy.

[0087]

[0088] Table 1

[0089]

[0090] Table 2

[0091]

[0092]

[0093] Table 3

[0094]

[0095] Table 4

[0096] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and not to limit the technical solutions. Those skilled in the art should understand that any modifications or equivalent substitutions to the technical solutions of the present invention without departing from the spirit and scope of the present invention should be covered within the scope of the claims of the present invention.

Claims

1. A medical image segmentation device based on a lightweight attention module, characterized in that: A segmentation model is built based on the UNet network. The segmentation model includes 4-6 upsampling and downsampling stages. Each upsampling and downsampling stage includes a convolutional sampling module. The downsampling layer in the upsampling and downsampling stages uses pooling operations, and the upsampling layer uses interpolation operations. Each convolutional layer is followed by a batch normalization layer. In the downsampling layer, a local spatial attention module formed by a combination of large-kernel dilated convolutions and a channel attention module for interaction between channels are added in series. The local spatial attention module is used to generate attention coefficients for different features in the local convolutional space, and the channel attention module is used to compress information along the channel direction and perform feature learning in the channel dimension. Convolutional gated attention modules are added to the skip connections and upsampling layers. The gated attention modules are used to filter irrelevant information and ensure the effectiveness of feature upsampling. A global attention module is set at the network bottleneck. The global attention module is used to calculate the correlation between each spatial location and all other spatial locations globally to realize the long-range dependency between pixel locations. Local Spatial Attention Module Use Convolution and GELU activation functions are used to activate the original feature parameters, followed by... Three large-kernel dilated convolutions of different sizes are used to extract image features at multiple scales, which are then fused to generate an attention weight map. Finally, a residual structure is added to connect the input to the result of the dot product of the attention map. The specific steps are as follows: （1） （2） （3） Where x represents the input feature information. For large-kernel dilated convolution, Conv is ordinary convolution. For attention weights, GELU is the activation function; The channel attention module uses average pooling to generate channel features, replacing depthwise convolutions with fully connected layers. Leveraging their highly non-linear global effects, it adaptively generates channel attention maps, as shown in the following formula: in represents the sigmoid activation function, Linear represents a fully connected structure, ReLU represents the activation function, and AvePool represents the average pooling operation; The gated attention module extracts information from a coarse scale as approximate features of the input tensor. This information is then multiplied by the detailed features extracted at a finer scale after being activated by different activation functions. This process performs a secondary feature selection, giving the final weighted map a detailed classification capability and accurately identifying the usefulness of information. The formula for the attention module is shown below: in The feature weight matrix, For sigmoid activation function, ; The global attention module multiplies the input features by their transpose matrix. Each point in the attention map is represented as the sum of the products of all channels at two spatial locations. Following the steps of matrix multiplication, the correlation between the spatial domain represented by each row and the global information can be calculated. The specific formula is as follows: in The position to be calculated It is the location to be used for weighting, that is, the location of the feature map; It is input. It is the output of the dot product. It is a function for calculating similarity. It is a single variable, and For normalization; for the original feature map First use Convolution to generate Then through Obtain similarity , Execute another , .

2. The medical image segmentation device based on a lightweight attention module according to claim 1, characterized in that: The segmentation model includes five upsampling and downsampling stages.

3. A training method for a medical image segmentation device, characterized in that: Includes the following steps: (1) Obtain the image set, normalize the content of all images in the image set to a value with an average value of 0 and a standard value of 1, and adjust the resolution to 256×256. Then divide the image set into training set and validation set. (2) Set the hyperparameters and loss function, initialize the optimizer, and initialize the parameters used for image segmentation. Neural networks The neural network is the segmentation model based on the UNet network as described in any one of claims 1 or 2; (3) Input the training set x The neural network obtains the segmentation result y, calculates the loss between the predicted result and the true value, and then backpropagates to update the network parameters; (4) Use the validation set to update the network parameters in step (3). The neural network is used for validation. The Dice coefficient and IOU of the current model on the validation set are calculated. If the Dice coefficient is better than the historical maximum, the current model parameters are saved, and the historical best Dice value is updated. ； (5) Training ends when the number of iterations reaches the set value.

4. The training method for the medical image segmentation device according to claim 3, characterized in that: The loss function uses a combination of cross-entropy and dice loss, and its formula is as follows: Where N is the total number of samples. , , .

Images