A lightweight remote sensing image semantic segmentation method and device based on knowledge distillation

Abstract

The application discloses a kind of light weight remote sensing image semantic segmentation methods based on knowledge distillation, comprising: constructing the network including one extraction local spatial feature branch CNN, one extraction global spatial feature branch Transformer, the feature fusion module of fusion local spatial feature and global spatial feature and the decoder of the decoding of fusion feature;ResNet50 and SegFormer-B3 are replaced CNN and Transformer, and the teacher network is obtained;ResNet18 and SegFormer-B0 are replaced CNN and Transformer, and the student network is obtained;Dynamic temperature-based fusion feature distillation module is constructed to jointly train student network and teacher network, dynamically search temperature parameters, using structured distillation method at feature level, so that student network pays attention to the feature representation of each channel in the fusion feature of teacher network during training, and a light weight remote sensing image semantic segmentation model based on knowledge distillation is obtained.The scheme of the application reduces the precision loss of remote sensing image semantic segmentation, and realizes model light weight.

Description

Technical Field

[0001] This invention relates to the field of remote sensing image information processing technology, and in particular to a lightweight remote sensing image semantic segmentation method and apparatus based on knowledge distillation. Background Technology

[0002] Semantic segmentation of remote sensing images refers to labeling each pixel in a remote sensing image according to its category, thereby enabling the identification and analysis of features, targets, scenes, etc., within the image. Semantic segmentation is one of the fundamental and key tasks of remote sensing image understanding, with wide-ranging applications such as urban planning, land use, environmental monitoring, and military reconnaissance.

[0003] In recent years, deep learning technology has made significant progress in the field of semantic segmentation of remote sensing images, especially methods based on convolutional neural networks (CNNs). CNNs have powerful feature extraction capabilities, enabling them to extract rich local spatial information from remote sensing images, thereby improving the accuracy of semantic segmentation. However, due to the limitation of their receptive field, CNNs have limitations in modeling long-range contextual information.

[0004] To overcome the limitations of CNNs, some research has begun to introduce the Transformer architecture to enhance the feature representation capabilities of deep learning models. The Transformer architecture is an encoder-decoder architecture based on a self-attention mechanism, capable of effectively processing sequential data, such as text data in natural language processing. The Transformer architecture has been successfully applied in computer vision, such as image classification, object detection, and image generation.

[0005] However, applying the Transformer architecture to remote sensing image semantic segmentation tasks still faces some challenges. On the one hand, how to fully utilize the feature extraction capabilities of CNNs and Transformers to enhance the accuracy of remote sensing image semantic segmentation is a key issue. In addition, for scenarios such as on-orbit processing of remote sensing satellites and airborne processing of UAVs, lightweighting the model is also an important issue when carrying out real-time semantic segmentation of remote sensing images on edge devices with limited storage and computing resources. Summary of the Invention

[0006] To address the aforementioned problems, this invention provides a lightweight remote sensing image semantic segmentation method and apparatus based on knowledge distillation, wherein the method includes the following steps:

[0007] S1. Acquire remote sensing images and preprocess them, then divide the preprocessed remote sensing images into training and testing sets.

[0008] S2. Construct a semantic segmentation network for remote sensing images based on CNN-Transformer feature fusion, including a CNN branch for extracting local spatial features, a Transformer branch for extracting global spatial features, a feature fusion module for fusing the extracted local and global spatial features, and a decoder FPNHead for decoding the fused features.

[0009] S3. Replace the CNN and Transformer in S2 with ResNet50 and SegFormer-B3 to obtain the teacher network;

[0010] S4. The teacher network is trained using the training set in S1 to obtain a remote sensing image semantic segmentation model, and the accuracy of the model is evaluated using the test set in S1.

[0011] S5. Replace the CNN and Transformer in S2 with ResNet18 and SegFormer-B0 to obtain the student network;

[0012] S6. Construct a dynamic temperature-based fusion feature distillation module. Use the training set in S1 to train the student network. During training, dynamically search for temperature parameters. Use the structured distillation method at the feature level to enable the student network to focus on the channel dimension feature representation in the teacher network's fusion features during training. This will result in a lightweight remote sensing image semantic segmentation model based on knowledge distillation.

[0013] S7. Based on the test set in S1, evaluate the accuracy of the lightweight remote sensing image semantic segmentation model based on knowledge distillation, and use the lightweight remote sensing image semantic segmentation model based on knowledge distillation to perform fast semantic segmentation on newly acquired remote sensing images.

[0014] This invention also proposes a lightweight remote sensing image semantic segmentation device based on knowledge distillation, comprising:

[0015] processor;

[0016] A memory on which computer programs that can run on the processor are stored;

[0017] When the computer program is executed by the processor, it implements a lightweight remote sensing image semantic segmentation method based on knowledge distillation.

[0018] The beneficial effects of the technical solution provided by this invention are:

[0019] This invention inputs training samples into a teacher network and a student network to obtain their respective encoded features and semantic segmentation results. It obtains the temperature parameter for knowledge distillation through gradient descent. The invention calculates the channel loss between the encoded features output by the teacher and student networks at different scales, the KL divergence between the semantic segmentation results output by the teacher and student networks, and the cross-entropy loss function between the ground truth and the semantic segmentation results output by the student network. These three factors are then weighted and summed to obtain a total loss function. Backpropagation and parameter updates are performed on the student network based on this total loss function, thereby achieving knowledge distillation. Using this invention, a lightweight model is achieved while reducing the loss of accuracy in remote sensing image semantic segmentation. Attached Figure Description

[0020] Figure 1 This is a flowchart illustrating a lightweight remote sensing image semantic segmentation method based on knowledge distillation, as described in this invention.

[0021] Figure 2 This is a structural diagram of the CNN-Transformer feature fusion network used for semantic segmentation of remote sensing images in an embodiment of the present invention;

[0022] Figure 3 This is a structural diagram of the Transformer Block according to an embodiment of the present invention;

[0023] Figure 4 This is a structural diagram of Bottlenect, an embodiment of the present invention;

[0024] Figure 5 This is a structural diagram of the feature fusion module in an embodiment of the present invention;

[0025] Figure 6 This is a network structure diagram of the teacher network according to an embodiment of the present invention;

[0026] Figure 7 This is a network structure diagram of the student network according to an embodiment of the present invention;

[0027] Figure 8 This is a schematic diagram of a lightweight remote sensing image semantic segmentation method based on knowledge distillation, according to an embodiment of the present invention.

[0028] Figure 9 This is a structural diagram of the fusion feature distillation module based on dynamic temperature according to an embodiment of the present invention. Detailed Implementation

[0029] To make the objectives, technical solutions, and advantages of the present invention clearer, the embodiments of the present invention will be further described below with reference to the accompanying drawings.

[0030] A flowchart of a lightweight remote sensing image semantic segmentation method based on knowledge distillation in this embodiment is shown below. Figure 1Specifically, it includes the following steps:

[0031] S1: Acquire remote sensing images, perform image enhancement operations such as rotation, scaling, flipping, and random photometric distortion on the remote sensing images to obtain a dataset, and divide it into training set and test set.

[0032] S2. Construct a remote sensing image semantic segmentation network based on CNN-Transformer feature fusion. This network consists of an encoder and a decoder.

[0033] The encoder includes a two-branch global-local feature extraction structure based on Transformer and CNN, as well as a feature fusion module. The two-branch global-local feature extraction structure consists of two sub-networks: a global feature extraction sub-network based on Transformer and a local feature extraction sub-network based on CNN.

[0034] In a further embodiment:

[0035] refer to Figure 2 The Transformer subnetwork uses a multi-layer Transformer encoder to extract long-range contextual information from the input remote sensing image. The multi-layer Transformer encoder consists of a patch embedding module for segmenting the input image into 4×4 resolution and four Transformer blocks. The input dataset first passes through a patch embedding, then through the first Transformer block to obtain the first global feature T1, then through the second Transformer block to obtain the second global feature T2, the third Transformer block to obtain the third global feature T3, and the fourth Transformer block to obtain the fourth global feature T4. After feature extraction by the four Transformer blocks, four feature maps with resolutions of [H / 4×W / 4, H / 8×W / 8, H / 16×W / 16, H / 32×W / 32] are finally obtained.

[0036] Each Transformer block utilizes a self-attention mechanism to extract deep global semantic features from the input features, reducing the feature map resolution to half that of the input Transformer block features. Each Transformer block contains multiple self-attention calculation modules and a hybrid feedforward neural network (Mix-FFN), as well as a patch merging module. (Reference) Figure 3The Transformer Block consists of N modules, each consisting of a Self-Attention Calculation module and a MIX-FNN module connected together, and then connected to a PatchMerging module.

[0037] refer to Figure 2 The CNN subnetwork uses a multi-layer CNN encoder to extract local spatial information from the input remote sensing image. The input sequentially passes through one Stem, one Maxpool, and the first Bottlenect to obtain the first local spatial feature C1; the second Bottlenect to obtain the second local spatial feature C2; ​​the third Bottlenect to obtain the third local spatial feature C3; and the fourth Bottlenect to obtain the fourth local spatial feature C4. Each Bottlenect group consists of N Bottlenects.

[0038] refer to Figure 4 Bottlenect consists of the sum of the outputs of two branches followed by a ReLU. The two branches have the same input. One branch consists of one Conv1×1 and one Norm; the other branch consists of two Conv1×1, one Conv3×3, two Norm&ReLU, and one Norm. The input passes through the first Conv1×1, the first Norm&ReLU, Conv3×3, the second Norm&ReLU, the second Conv1×1, and the Norm before the output. Norm&ReLU represents the combination of the normalization function and the activation function.

[0039] ResNet can replace CNN-Transformer. ResNet consists of several residual blocks called bottlenecks, each containing a residual structure. The residual structure uses short connections to transfer features from shallow to deep layers, solving the gradient vanishing problem in deep CNN networks. The formula is H(x) = F(x) + x, where F(x) is the output obtained after passing through convolutional layers and activation functions. Based on the combination of the number of bottlenecks in layers 1-4, ResNet can be divided into ResNet18, ResNet34, ResNet50, ResNet101, and ResNet152. For example, the number of bottlenecks in layers 1-4 of ResNet50 is [3, 4, 6, 3].

[0040] refer to Figure 2The feature fusion module fuses features extracted at different scales to obtain fused features with global and local information. The remote sensing image semantic segmentation network based on CNN-Transformer feature fusion has four feature fusion modules. The first feature fusion module fuses features C1 and T1 to form fused feature F1; the second feature fusion module fuses features C2 and T2 to form fused feature F2; the third feature fusion module fuses features C3 and T3 to form fused feature F3; and the fourth feature fusion module fuses features C4 and T4 to form fused feature F4.

[0041] refer to Figure 5 , Figure 5 This is a structural diagram of the feature fusion module used for semantic segmentation of remote sensing images in this invention. Local spatial features and global features are the inputs to the feature fusion module. The local spatial features and global features are first processed through their respective Conv1×1 operations, then mixed and sequentially passed through a common Concat, the first Conv1×1, Norm&Act, and the second Conv1×1, before being added to the output of the second branch. The input of the second branch is a mixture of local spatial features and global features. The second branch consists of one Concat and one Conv1×1 operation. The input of the second branch is processed through Concat and Conv1×1 before being output.

[0042] The multi-scale feature fusion module is designed to fuse local features from the CNN branch and global features from the Transformer branch, enabling efficient and flexible fusion of features with different resolutions and channel numbers. In the feature fusion module, the features extracted by the CNN and Transformer branches (named C and T) are first adjusted in dimension using convolutional networks with 1×1 kernels, and then C and T are concatenated and merged. At this stage, the concatenated and merged features are separated into a third dimension. The merged features are then fed into a convolutional network with 1×1 kernels for normalization and activation. Finally, the features from the CNN and Transformer branches are fused through a single convolutional layer with 1×1 kernels. Furthermore, the feature fusion module incorporates a residual structure to accelerate model convergence. The entire network contains four feature fusion modules, accepting CNN and Transformer features with spatial resolutions of H / 4×W / 4, H / 8×W / 8, H / 16×W / 16, and H / 32×W / 32, respectively.

[0043] refer to Figure 2The decoder uses the FPNhead from the Feature Pyramid Network as its decoding head. It upsamples and further fuses the multi-scale fused features output by the encoder to obtain a semantic segmentation result of the same size as the input image. F1 takes a Conv3×3 & ReLU input as its output, F2 takes one Scale Head input as its output, F3 takes two Scale Head inputs as its output, and F4 takes three Scale Head inputs as its output. These four outputs are then superimposed and input into a Conv1×1 module for final output. Each Scale Head consists of one Conv3×3, one ReLU, and one upsampling module, which are sequentially connected. Conv3×3 & ReLU represents a combination of a 3×3 convolution and an activation function.

[0044] S3. Replace the CNN and Transformer in S2 with ResNet50 and SegFormer-B3 to obtain the teacher network. Refer to the network structure diagram of the teacher network. Figure 6 .

[0045] S4. The teacher network is trained using the training set in S1 to obtain a remote sensing image semantic segmentation model, and the accuracy of the model is evaluated using the test set in S1.

[0046] The training loss function for step S4 is:

[0047]

[0048] Where N is the number of pixels in the remote sensing image, M is the number of categories, and y n,m It is the true category label of the nth pixel and the mth class, p n,m It is the probability that n pixels are of class m as predicted by the semantic segmentation model of remote sensing images, and ln represents the natural logarithm operation.

[0049] S5. Replace the CNN and Transformer in S2 with ResNet18 and SegFormer-B0 to obtain the student network. Refer to the network structure diagram of the student network. Figure 7 .

[0050] S6. Construct a dynamic temperature-based fusion feature distillation module and use the training set in S1 to train the student network. This module can dynamically search for temperature parameters during training and use a structured distillation method at the feature level, enabling the student network to focus on the channel dimension feature representation in the teacher network's fusion features during training, resulting in a lightweight remote sensing image semantic segmentation model based on knowledge distillation.

[0051] Figure 8This is a schematic diagram of the lightweight remote sensing image semantic segmentation method based on knowledge distillation of the present invention, which includes four parts: teacher network, student network, loss function, and dynamic temperature-based fusion feature distillation module. T F2 T F3 T F4 T These represent the fused features extracted by the encoder of the teacher network, obtained by fusing the features extracted by the CNN branch (ResNet50) and the Transformer branch (SegFormer-B3) of the encoder. Similarly, F1... S F2 S F3 S F4 S The feature extracted by the CNN branch (ResNet18) and the Transformer branch (SegFormer-B0) in the student network encoder is obtained by feature fusion. T and M S The 'H' and 'W' represent the classification probability distribution of each pixel output after the fusion features of the teacher network and student network are decoded by the feature decoder, respectively, with dimensions H×W×N. Here, (H,W) represents the image resolution, and N is the number of categories in the semantic segmentation task. The dashed line represents the backpropagation process.

[0052] Figure 9 This is a structural diagram of the dynamic temperature-based fusion feature distillation module. This module can efficiently transfer feature-level knowledge from the teacher model to the student model, thereby better utilizing rich multi-scale features and improving the segmentation accuracy and robustness of the student model. The dynamic temperature-based fusion feature distillation module designed in this invention comprises two components: a dynamic temperature generation module and a feature channel-based knowledge distillation loss function.

[0053] The core of dynamic temperature is to set the hyperparameter T in knowledge distillation as a learnable network parameter θ. temp During training, the learning difficulty of the student model is controlled by gradually increasing the distillation loss between the student and teacher models. The optimization objective is:

[0054]

[0055] Among them, f t (x;θ tea ) indicates that the teacher network parameter is θ tea The output probability distribution when the input is x; f s (x;θ stu ) represents the student network parameters as θ stu The output probability distribution is given by input x, where y is the true class label.

[0056] Applying the alternation algorithm, which involves fixing one set of variables and solving for another set of variables to solve the above formula, we can transform it into:

[0057]

[0058]

[0059] The optimization process described above can be performed using stochastic gradient descent (SGD) with a learning rate of μ, and its formula is as follows:

[0060]

[0061] In practice, the dynamic distillation described above is achieved through a non-parametric gradient reversal layer (GRL), as shown in the attached figure. Figure 8 As shown by the dashed line in the image.

[0062] The knowledge distillation loss function based on feature channels can be expressed as:

[0063]

[0064] Among them, L dffd F represents the feature loss calculated in the dynamic temperature-based fusion feature distillation module. i T ,i=1,2,3,4 represent the features output by the i-th feature fusion module of the teacher network, F i S Let i = 1, 2, 3, 4 represent the features output by the i-th feature fusion module of the student network, H represent the feature map height, W represent the feature map width, c = 1, 2, ..., C represent the number of channels in the feature map, and j = 1, 2, ..., H·W represent the number of pixels. The feature output by the i-th feature fusion module of the teacher network at the j-th pixel and c-th channel. For the feature output by the i-th feature fusion module of the student network at the j-th pixel and c-th channel, T = θ temp , is the distillation temperature parameter, Softmax is the activation function, t i Let s represent the output probability distribution of the teacher network. i Let y represent the output probability distribution of the student network. i This represents the actual category label.

[0065] In summary, the knowledge distillation process of this invention includes the following steps:

[0066] (1) For each training sample, input it into the teacher network and the student network to obtain their respective encoding features and semantic segmentation results.

[0067] (2) Obtaining the temperature parameter θ of knowledge distillation through gradient. temp .

[0068] (3) Calculate the channel loss L between the encoded features output by the teacher network and the student network at different scales. dffd .

[0069] (4) Calculate the KL divergence loss function L between the semantic segmentation results output by the teacher network and the student network. kl The cross-entropy loss function L between the ground truth and the semantic segmentation results output by the student network. ce The calculation formulas are as follows:

[0070]

[0071] L ce =-y i ln(s i )

[0072] Where T is the distillation temperature parameter, Softmax is the activation function, and t i Let s represent the output probability distribution of the teacher network. i Let y represent the output probability distribution of the student network. i This represents the actual category label.

[0073] (5) L dffd L kl and L ce The weighted sum of these three factors yields a total loss function L = αL dffd +βL kl +γL ce , where α, β, and γ are the weights of the three feature components, which are manually adjusted hyperparameters. Backpropagation and parameter updates are performed on the student network based on the total loss function L, thereby achieving knowledge distillation. Steps (1) to (4) are repeated until the student network converges, finally obtaining the trained lightweight model.

[0074] S7. Based on the test set in S1, evaluate the accuracy of the trained lightweight remote sensing image semantic segmentation model based on knowledge distillation, and use the lightweight remote sensing image semantic segmentation model based on knowledge distillation to perform fast semantic segmentation on newly acquired remote sensing images.

[0075] The embodiment also includes a lightweight remote sensing image semantic segmentation device based on knowledge distillation, comprising:

[0076] processor;

[0077] A memory on which computer programs that can run on the processor are stored;

[0078] Among them, when the computer program is executed by the processor, it implements a lightweight remote sensing image semantic segmentation method based on knowledge distillation.

[0079] The above description of the disclosed embodiments enables those skilled in the art to make or use the invention. Various modifications to these embodiments will be readily apparent to those skilled in the art, and the general principles defined herein may be implemented in other embodiments without departing from the spirit or scope of the invention. Therefore, the invention is not to be limited to the embodiments shown herein, but is to be accorded the widest scope consistent with the principles and novel features disclosed herein.

Claims

1. A lightweight semantic segmentation method for remote sensing images based on knowledge distillation, characterized in that, Includes the following steps: S1. Acquire remote sensing images and preprocess them, then divide the preprocessed remote sensing images into training and testing sets. S2. Construct a remote sensing image semantic segmentation network based on CNN-Transformer feature fusion, including a CNN branch for extracting local spatial features, a Transformer branch for extracting global spatial features, a feature fusion module for fusing the extracted local spatial features and global spatial features, and a decoder FPNHead for decoding the fused features. S3. Replace the CNN and Transformer in S2 with ResNet50 and SegFormer-B3 to obtain the teacher network; S4. The teacher network is trained using the training set in S1 to obtain a remote sensing image semantic segmentation model, and the accuracy of the model is evaluated using the test set in S1. S5. Replace the CNN and Transformer in S2 with ResNet18 and SegFormer-B0 to obtain the student network; S6. Construct a dynamic temperature-based fusion feature distillation module. Use the training set in S1 to train the student network. Based on the remote sensing image semantic segmentation model, dynamically search for temperature parameters during training. Use the structured distillation method at the feature level to enable the student network to focus on the channel dimension feature representation in the teacher network's fusion features during training, thus obtaining a lightweight remote sensing image semantic segmentation model based on knowledge distillation. S7. Based on the test set in S1, evaluate the accuracy of the lightweight remote sensing image semantic segmentation model based on knowledge distillation, and use the lightweight remote sensing image semantic segmentation model based on knowledge distillation to perform fast semantic segmentation on newly acquired remote sensing images. The training loss function in step S6 is: in, They are The weights are manually adjusted hyperparameters. This represents the feature loss calculated in the dynamic temperature-based fusion feature distillation module, used to measure the inter-channel similarity between fused features. This represents the KL divergence loss function, used to measure the difference between the probability distributions of the student network and the teacher network outputs. It is the cross-entropy loss between the output probability distribution of the student network and the true class label; Based on the total loss function Backpropagation and parameter updates are performed on the student network to achieve knowledge distillation until the student network converges, and finally a lightweight remote sensing image semantic segmentation model based on knowledge distillation is obtained after training. In step S6, the dynamic temperature is used to set the hyperparameter T in knowledge distillation to a learnable network parameter. During training, the learning difficulty of the student network is controlled by gradually increasing the distillation loss between the student network and the teacher network. The optimization objective is: in, Indicates the teacher network parameters are Input is The output probability distribution at that time; Indicates the student network parameters are Input is The output probability distribution at that time, These are real category labels; Learning rate The stochastic gradient descent algorithm for Update: Where u is the learning rate.

2. The lightweight remote sensing image semantic segmentation method based on knowledge distillation according to claim 1, characterized in that, Step S1 involves preprocessing the remote sensing image, including rotating, scaling, flipping, and applying random photometric distortion to the image.

3. The lightweight remote sensing image semantic segmentation method based on knowledge distillation according to claim 1, characterized in that, The number of feature fusion modules for both the teacher network and the student network is 4.

4. The lightweight remote sensing image semantic segmentation method based on knowledge distillation according to claim 1, characterized in that, The training loss function for step S4 is: Where N is the number of pixels in the remote sensing image, and M is the number of categories. It is the true category label of the nth pixel and the mth class. It is predicted by the semantic segmentation model of remote sensing images. The pixel is the first The probability of a class, where ln represents the natural logarithm.

5. A lightweight remote sensing image semantic segmentation method based on knowledge distillation according to claim 1, characterized in that, , , Specifically, it is expressed as follows: in, Indicates the teacher network The features output by the feature fusion module Indicates the student network The features output by the feature fusion module are shown, where H represents the feature map height and W represents the feature map width. This represents the number of channels in the feature map. Indicates the number of pixels. For Teachers' Network The output of the feature fusion module at the th The pixel, the first Characteristics of the channel For student network The output of the feature fusion module at the th The pixel, the first Characteristics of the channel , where is the distillation temperature parameter, and Softmax is the activation function. This represents the output probability distribution of the teacher network. This represents the output probability distribution of the student network. This represents the actual category label.

6. A lightweight remote sensing image semantic segmentation device based on knowledge distillation, characterized in that, The device includes: processor; A memory on which computer programs that can run on the processor are stored; When the computer program is executed by the processor, it implements a lightweight remote sensing image semantic segmentation method based on knowledge distillation as described in any one of claims 1 to 5.

Images