Construction method of lightweight remote sensing target instance segmentation model based on state space model knowledge distillation

Abstract

The invention provides a method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation. The method comprises the following steps of: 1, reading data by a computer; the method comprises the following steps of: reading an unlabeled SA-1B data set and labeled remote sensing image data by using a computer, and respectively preprocessing the image data; 2, constructing a knowledge distillation network based on the state space model and a lightweight remote sensing target instance segmentation network; and step 3, training the neural network. After a network is constructed by using a Pytorch deep learning framework, a backbone model is trained by using unlabeled data, and then the complete network is trained by using labeled remote sensing data. And step 4, remote sensing image instance segmentation. And high-quality remote sensing image instance segmentation can be realized on a test set or an actual image by using the network model trained in the step 3. Compared with a traditional method, model parameters can be greatly reduced, model reasoning time is saved, and the method is more suitable for various scenes needing real-time application or quick response.

Description

Technical Field

[0001] The invention relates to a method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation, and belongs to the field of satellite networks. Background Art

[0002] The state space model (SSM) is a model that describes the previous state of the model and predicts its next state based on the current input. SSM was first applied in the control field, taking a continuous sequence as input and predicting the output sequence. The model has two equations, the state equation and the output equation. The state equation is used to describe the result of multiplying the matrix by the input and the matrix by the previous state. The output equation describes how the state is converted to the output through the matrix, and how the input affects the output through the matrix.

[0003] In order to solve the quadratic complexity of the attention mechanism in the Transformer architecture and the sequence forgetting problem of RNN at the same time, and to obtain a model with the advantages of both, Albert Gu et al. designed the S4 model. The original SSM model can be discretized to use discrete sequences as input. The S4 model introduces the memory mechanism of time step and HiPPO. Every time a discrete signal is received, the signal value is retained until a new discrete signal is received. The processing interval of the operation is a new learnable parameter, called the time step. In this way, the model is maintained in stages, and the input signal is sampled according to the learned time step to generate the final output. There are two ways of discretization operation: zero-order hold and Euler discretization, but the latter is used more. At the same time, the discretized state space model has a convolution-like expression. In the classic image recognition task, the convolution kernel is used to aggregate features, and the expression can be expanded to be written in a form similar to convolution, so that the convolution mode can be used for training, and the cyclic mode reasoning can be used for testing, which greatly improves the representation efficiency.

[0004] For long-distance problems, HiPPO is introduced, which is also a concept proposed by Albert Gu et al. It compresses all the input signals currently seen into a coefficient vector, which can well capture the recent token and decay the old token state representation, and generate the optimal solution of the state matrix through function approximation.

[0005] Based on the work of their predecessors, Albert Gu and others proposed the Mamba model, which is a culmination of their predecessors' achievements. Mamba was inspired by the Simplified State Space Layers for Sequence Modeling (S5), and adopted a hardware-aware algorithm to reduce IO access between computer layers. It also introduced the selective information processing of structured SSM (S6 for short), whose underlying logic is represented as a learnable input mapping function, which will produce different calculation parameters for different inputs. At the same time, Mamba also borrowed the concept of LSTM and gate circuits, merging the SSM structure with the gated MLP block of Transformer into one layer to simplify the past deep sequence model architecture. In terms of task performance, Mamba has achieved performance that exceeds SOTA in language, audio, and DNA sequence modalities, and surpasses Transformers of the same size in language tasks, and is close to Transformers that are twice as large.

[0006] After the Mamba model was proposed, it attracted widespread attention and was extended to other directions by many scholars. The Vision Mamba model (ViM) came from this. The ViM model is similar to the Vision Transformer (ViT). It first divides the input image into multiple small patches, then flattens the small patches into sequences and maps them into tokens. After that, it is uniformly input into the stacked layers of ViM after position encoding and adding category tensors. In each layer, it undergoes two different linear mappings, and processes the sequence forward and backward in the middle. Finally, the Sigmoid Linear Unit (SiLU) activation function is used. It should be noted that for non-classification tasks, no additional learnable classification tags are required. Finally, ViM has better results and efficiency in benchmarks such as ImageNet classification, COCO object detection, and ADE20K semantic segmentation, and for high-resolution images, it is 2.8 times faster than DEIT while reducing 86% GPU memory.

[0007] In recent months, Mamba can be roughly divided into Phase 1 and Phase 2 according to the development of the general backbone network. The original Mamba proposed by Albert Gu et al. introduced the concepts of selective state space model and stackable state space block (State Space Model Block), realized the fusion of S6 and gated multilayer perceptron (MultilayerPerceptrons, MLP), and provided an open source network layer for use; inspired by it, Vision Mamba and VMamba appeared three months later, both of which improved it from different angles and became the two main improvement ideas for subsequent researchers. The former was proposed by L. Zhu et al. in January 2024. It combines the idea of ​​Transformer, replaces the stacking layer of Transformer with the original Mamba layer, and uses image slicing and encoding to process, realizing Vision Mamba with good visual effects; while the latter was proposed by Y. Liu et al. almost at the same time. The main consideration is that Mamba lacks the viewing order of image slicing. The original slicing method may cause the images that are close to each other to be separated farther apart. Therefore, they proposed the idea of ​​fast scanning, using SS2D Block to perform multiple scanning methods and cooperating with the internal VSSBlock for stacking operations. It should be noted that VSS Block and Vim Block are similar because they are both developed on the basis of Mamba, but the core difference is that VSS Block does not have the bidirectional scanning structure inside Mamba, which is caused by the fact that he puts the scanning order in the front SS2D Block. In the future, people proposed Mamba Mixer for visual Mamba, and Mamba2 for traditional Mamba, and constructed the mathematical connection between Mamba and Transformer.

[0008] Model lightweighting (Distributed computing) refers to a method of obtaining a more efficient model by reducing model storage and speeding up inference speed without affecting model performance within the allowable range through lightweight methods. It can play a good practical role in many practical application scenarios. For example, in satellite remote sensing scenarios, remote sensing satellites have been used in a variety of scenarios. The new low-orbit giant star constellation network makes up for the shortcomings of traditional high-orbit satellites in resolution or real-time detection, and has also become a research hotspot in various countries. However, in remote sensing target detection, especially in low-orbit satellite detection, there is a contradiction between real-time performance and speed accuracy. Satellite space resources are limited, and complex deep learning models are difficult to deploy. Simple models are difficult to cope with complex remote sensing scenarios. Therefore, it is crucial to design lightweight target detection algorithms suitable for remote sensing scenarios.

[0009] Next, looking back at the field of object detection over the past 20 years, whether it is the single-stage and two-stage methods produced 12 years later, or the vision transformer introduced from the NLP field 20 years ago, people have always been pursuing faster and better methods. In recent years, the basic visual model, also known as the large model, has achieved excellent performance and good scalability and generalization in various visual tasks through a substantial expansion of data and model scale, and has become a major research hotspot. However, although the performance of the large visual model is good, the inference speed is slow and the number of parameters is too large. For example, the large SAM model for object segmentation uses more than 1 billion parameters. The goal is to use a lightweight method to provide a basis for the deployment of basic models on satellites while ensuring accuracy and significantly compressing parameters and power consumption.

[0010] There are three specific methods for model lightweighting: knowledge distillation that transfers knowledge from a large teacher model to a small student model, model pruning with network sparsification as the core idea, and model quantization that converts floating-point algorithms to fixed-point algorithms from a storage perspective. However, pruning cannot fundamentally accelerate model reasoning, and quantization generally cannot achieve an exaggerated improvement, so we start with the perspective of knowledge distillation.

[0011] The present invention proposes a method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation. By constructing a backbone network with strong learning ability according to the state space model, using knowledge distillation and using ViT as a teacher model to learn the complex knowledge of a large network, the expandability of the state space network layer is used to quickly calculate the results of multi-layer stacking to speed up the reasoning speed, and the constructed student network itself greatly reduces the storage, and the neck network theoretically alleviates the loss of small targets to a certain extent. The time of the total process of model reasoning and the space occupied by the model storage are greatly reduced, and the time cost and space cost required for the model operation can be greatly reduced. A practical solution is proposed for the contradiction between limited satellite resources and high model performance requirements. The present invention aims at aerospace remote sensing problems, constructs a set of lightweight construction methods for remote sensing instance segmentation models based on state space model knowledge distillation, which can segment masks while detecting targets, and provides a reference configuration method. In addition, this technology can also be easily migrated to other mission scenarios that require lightweight, such as vehicle-mounted detection, agricultural monitoring, and aerial ground detection. Summary of the invention

[0012] The purpose of the present invention is to design a method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation. This method can make the network model faster and have smaller storage than the conventional model without distillation, and can be efficiently trained, tested and used on a conventional graphics processing unit (GPU), and achieve the same effect as the conventional model. Specifically, the present invention is implemented by the following technical solutions:

[0013] The present invention is a method for constructing a lightweight remote sensing target instance segmentation model based on state-space model knowledge distillation. It is essentially a "pre-training + fine-tuning" strategy. This method constructs a lightweight encoding backbone to learn the knowledge of a large model as pre-training weights; then connects the neck and the detection head, and sets all parameters to be trainable, and then fine-tunes the remote sensing data set to achieve the desired effect. The backbone network trained for the first time is essentially a high-dimensional feature extractor responsible for extracting multi-level image information; the neck network trained for the second time is used for multi-level feature fusion, and the head is responsible for processing specific tasks. The parameters of the backbone network are not frozen to ensure its specialization on remote sensing tasks to achieve better results. The specific steps of this method are as follows:

[0014] Step 1: Computer reads data. Use a computer to read the unlabeled SA-1B dataset and labeled remote sensing image data, and preprocess the image data respectively.

[0015] Step 2: Construct a knowledge distillation network based on the state-space model and a lightweight remote sensing target instance segmentation network.

[0016] The present invention proposes a construction method based on a state-space model knowledge distillation network and a lightweight remote sensing target instance segmentation network. The state-space model knowledge distillation network consists of two parts: a teacher network and a student network. The lightweight remote sensing target instance segmentation network consists of three parts: a backbone, a neck, and a head. Among them, the teacher network is an image encoder of the SAM basic model, using the parameters provided by the official, and no parameter adjustment is performed during training; the backbone network is the student network obtained after training based on the state-space model knowledge distillation network, which is responsible for extracting multi-level image information and is essentially a high-dimensional feature extractor; the neck network is used for multi-level feature fusion; and the head is responsible for processing specific tasks.

[0017] Step 3: Train the neural network. After constructing the state-space model-based knowledge distillation network using the Pytorch deep learning framework, use the unlabeled data to train this network, save the training parameter results, import the student network based on the state-space model knowledge distillation network, and use it as the backbone network of the lightweight remote sensing target instance segmentation network. After connecting the neck and head, use the labeled remote sensing data to train the lightweight remote sensing target instance segmentation network (i.e., the backbone, neck, and head network). The specific training method is: in the first training, the backbone is constructed through the state-space model, and then it is used as a student model to train the lightweight method of knowledge distillation to learn the complex knowledge of the large teacher model, and then the training results of the backbone are retained as pre-training weights, the neck and head are connected, and all parameters are fine-tuned on the remote sensing data set to ensure that it is specialized in remote sensing tasks, so that it can finally achieve the expected effect. In the training process, the optimal lightweight remote sensing target instance segmentation network parameters are selected according to the test data, and the lightweight remote sensing target instance segmentation network parameters at this time are saved as the final training results.

[0018] Step 4: Remote sensing image instance segmentation. Using the lightweight remote sensing target instance segmentation network model trained in step 3, high-quality remote sensing image instance segmentation can be achieved on the test set or actual image.

[0019] The advantages of the present invention are: the present invention uses the method of knowledge distillation to greatly reduce the number of parameters required by the model, and the selected state space model ensures the speed of reasoning and training. The "pre-training + fine-tuning" strategy adopted during training avoids the possibility of large deviations after feature fusion due to slight differences between the student model and the teacher model during learning. At the same time, according to specific tasks and objects, the required data sets can be selected for special training, ensuring the adaptability of the model under specific tasks and increasing its value in practical applications. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 It is the overall network structure diagram of the method of the present invention.

[0021] Figure 2 It is a diagram of the neck network structure of the method of the present invention.

[0022] Figure 3 It is a training structure diagram of the distillation of the method of the present invention.

[0023] Figure 4a-4f It is an image instance segmentation map of the method of the present invention. DETAILED DESCRIPTION

[0024] In order to better understand the technical solution of the present invention, the following Figure 1-Figure 4f The specific embodiments of the present invention are further described as follows:

[0025] The present invention is implemented in Python programming under the Pytorch framework. First, prepare the training and test data, then complete the network construction and configure the relevant parameters; then use the training data to train the network to obtain the optimal network parameters; finally use the trained network model parameters to test the test data to obtain the building extraction results. During the implementation process, the computer configuration is: the processor is Intel (R) Xeon (R) Platinum 8336C CPU, the main frequency is 3.50GHz, the memory is 251GB, the graphics card is NVIDIA GeForce RTX4090Ti, and the video memory is 24GB. The necessary configurations are Python3.10, Pytorch 2.1.2, Mmcv2.1.0, Mamba-ssm1.1.1, and causal-conv1d1.1.1.

[0026] Step 1: The computer reads the data. The data sources used in the present invention are the official SA-1B data set of SAM and the well-known public remote sensing data sets SSDD, WHU, and NWPU data sets. The SA-1B data set includes more than 1.1 billion segmentation masks collected on approximately 11 million licensed and privacy-protected images. About 12,000 images can be selected to basically meet the requirements. Basic data enhancement is performed on it and expanded to about 50,000 1024x1024 images. The SSDD, WHU, and NWPU data sets are all small and medium-sized data sets, and the three are respectively for optical SAR ships, static building targets, and mixed scenes of multiple types of targets. Select the complete data set and convert it to the COCO format in a unified manner to adapt to the MMCV framework. The data is also augmented and expanded to a size of 1024x1024 without changing the aspect ratio.

[0027] Step 2: Construct a knowledge distillation network based on the state-space model and a lightweight remote sensing target instance segmentation network.

[0028] Figure 1The overall description of the proposed method is shown. Knowledge distillation is a lightweight method. Unlike pruning and floating point fixed point storage quantization based on sparsification, distillation can extract knowledge from a bulky but effective large model and teach it to a small student model. The new model after distillation can have both the complexity of the original small model and the complex knowledge of the large model. In recent years, the state space model (SSM) has been introduced as a new method in computer vision, which has produced impressive performance in current perception, generation and other tasks, and its unique scalability when stacked has significantly improved the processing speed of feature extraction. The large model derived from the traditional vision transformer has achieved excellent performance in various visual tasks through a significant expansion of data and model scale, and has played a driving role in many current directions. However, although the performance of the large visual model is good, the inference speed is slow and the number of parameters is huge. For example, the large model of target segmentation SAM uses more than 1 billion parameters. The ultimate goal is to use lightweight methods to enable small models to have the performance of basic models while ensuring accuracy, greatly compressing parameters and power consumption, and finally deploy them in various required places.

[0029] The present invention proposes a construction method based on a state-space model knowledge distillation network and a lightweight remote sensing target instance segmentation network. The state-space model knowledge distillation network consists of two parts: a teacher network and a student network. The lightweight remote sensing target instance segmentation network consists of three parts: a backbone, a neck, and a head. Among them, the teacher network is an image encoder of the SAM basic model, using the parameters provided by the official, and no parameter adjustment is performed during training; the backbone network is the student network obtained after training based on the state-space model knowledge distillation network, which is responsible for extracting multi-level image information and is essentially a high-dimensional feature extractor; the neck network is used for multi-level feature fusion; and the head is responsible for processing specific tasks.

[0030] The backbone network is the core part of the present invention, which is responsible for extracting features from the original input. Figure 1 As shown. It consists of a series of stacked ViM block layers. The design goal of the backbone network is to use the global perception and memory of the state-space model to capture the semantic information and spatial position correlation of the input image and process the image information layer by layer. Similar to the previous article, the output of the backbone network is a four-dimensional feature map. For each batch, for each image in it, it is expanded to a size of h×h×c through clip() and divided into multiple p×p×c slices, containing a total of J=n 2Where h = p × c, h represents the length and width, c represents the dimension of the image, batch represents the training batch, p represents the side length of the small square, J represents the number of small squares, and the small square is also called patch. clip() is the official function of pytorch, which is used to stretch the image without changing the aspect ratio, enlarge (or reduce) it to the required size, and fill the insufficient part with black.

[0031] As attached Figure 1 As shown below, for the three different backbone designs, the Vanilla Mamba is composed of 12 layers of normal Vision Mamba Block (Vim Block), the model is the smallest, but the learning ability is weak; Trans Mamba is composed of 12 layers of Transformer Vision Mamba Block (TRVim Block) combined with Transformer, which is similar to the structure of Transformer, specifically including five components: Residual Layer, Normalization Layer, Mamba Layer, Gelu activation function layer and Multilayer Perceptrons (MLP), with strong learning ability, but the model is the largest and the reasoning is slow; Scanning Mamba is composed of 24 layers of Scanning Vision Mamba Block (SCVim Block), the model size is medium, but the learning ability is strong and the reasoning is fast. Users can choose according to task requirements, but the present invention recommends Scanning Mamba. The present invention is described in detail below from two aspects: intuitiveness and mathematical details.

[0032] The VanillaMamba is inspired by the original Mamba. It is simply composed of 12 layers of stacked normal Vision Mamba Blocks (Vim Block). However, in the training of the present invention, although it can converge in the first stage, its model may be relatively simple. At the same time, the present invention does not change the relevant parameters of the original memory. The graphics it processes are expanded from the original 256×256 size to 1024×1024, and the number of sequences is increased by 16 times, so its learning ability is relatively weak.

[0033] TransMamba has greatly increased parameters compared to the original backbone network, but correspondingly, efficiency has also been reduced to a certain extent. In addition, when the code is implemented, the present invention uses the SAM image encoder provided in the third-party library MMCV, and directly changes its stacking layer, resulting in a result that when the model is stored, a larger Python file will also be caused, because in the programming method, the original Transformer needs, and some parameters that the present invention does not need are also saved. When this model is doing a simple test, it is basically manifested as a bounding box (Bounding Box, BBox) that can barely converge, but the segmentation part performs poorly. In addition, the present invention does not replace the complete stacking layer by designing the visual Mamba block (Transformer Vision Mamba Block, TRVim Block) combined with Transfromer, but only replaces the attention layer therein, modifies the activation function at the same time, and modifies the loss function to the L1 norm, so as to achieve better results, because it may cause the value to be too small after taking the square, which is not conducive to parameter back transmission training.

[0034] Scanning Mamba focuses on the scanning method in model design. First, the stacking layer is transformed into the Mamba mixer provided by Facebook instead of the original SAM image encoder. At the same time, the number of stacked layers is increased to 24 to enhance learning ability. Secondly, the size of the embedding dimension is adjusted to reduce the dimension of the word vector and the parameters of the model. Finally, the scanning strategy of RS Mamba is selected, but instead of performing random scanning layers, only forward and backward scanning methods are performed. Specifically, when I input a series of word vectors, the order is reversed to obtain two different sequences, forward and backward. The parameters are shared when entering each layer. After the output, the reverse sequence is reversed again and thickened and averaged with the forward sequence, so that the output of this layer has the same dimension as the input, and then it is sent to the next layer.

[0035] For the J slices that have been obtained, use resize() to flatten them into a single dimension but with a length of p 2 ×c vectors, denoted as Where i represents the i-th image in the batch, 1≤i≤batch, k represents the vector corresponding to the k-th slice, and the arrangement order is scanned row by row. Each slice vector is multiplied by the trainable matrix W, which is a matrix with p columns. 2 c, a learnable matrix with d rows, and adding sine and cosine position encoding to each vector to obtain the matrix X0

[0036]

[0037] Among them, X position Learnable positional encodings.

[0038] Note that at this point, Vim's original formula is That is, the column vector needs to be consistent with the learnable category vector x class Front and back splicing, but no x is added here class This is because the task of the present invention is different from the original classification task. Secondly, in the downsampling stage of the neck, only an even number of slices can use two-dimensional convolution downsampling.

[0039] After that, X0 is input into the stack layer, Mamba_block. For VanillaMamba:

[0040] X b =X b-1 +Vim(X b-1 ),b=1,2,...,12

[0041] For TransMamba:

[0042] X b =X b-1 +TRVim(X b-1 ),b=1,2,...,12

[0043] For ScanningMamba:

[0044] X d =X d-1 +SCVim(X d-1 ),d=1,2,...,24

[0045] Among them, Vim() is the normal vision Mamba block mentioned above, TRVim() is the Transformer Vision Mamba Block (TRVim Block) mentioned above, SCVim() is the Scanning Vision Mamba Block (SCVim Block) mentioned above, X b refers to the output result of layer b, b is a positive integer between 1 and 12, d refers to the output result of the dth layer, The range of d is a positive integer between 1 and 24.

[0046] For f1(X i )=Vim(X i) is specifically handled as follows:

[0047] Normalize the output of the previous layer:

[0048] X' b-1 =norm(X b-1 )

[0049] Among them, X b-1 represents the output of the b-1 layer, norm() represents the normalization function, X' b-1 is the normalized result;

[0050] Through the linear() linear mapping function, the normalized result is mapped to the variable t b-1 and z b-1 ;

[0051] t b-1 = linear1(X' b-1 )

[0052] z b-1 = linear2(X' b-1 )

[0053] Among them, t b-1 Used to participate in normal forward and reverse operations, z b-1 Responsible for deciding how to merge the positive and negative results;

[0054] The parameter matrix B is then constructed by b-1 , C b-1 and Δ b-1 ;

[0055] t′ b-1 =Silu(conv1d(X' b-1 ))

[0056] B b-1 = linear3(t' b-1 )

[0057] C b-1 = linear4(t' b-1 )

[0058]

[0059] Among them, Silu() is the Sigmold activation function, linear p () are different linear mapping functions, p is a positive integer between 1 and 5, and the other symbols are interpreted as before;

[0060]

[0061] and is the discretized matrix, g() is the discretization function, X b The output of layer b is the intermediate variables X1, X3, ..., X23 of layers 1, 3, 5, ..., 23. 23 Enter the next layer, SSM() is the state space model formula constructed by it, and then substitute t′ obtained in the previous article b-1 Participate in the calculation and finally get the output y b-1 ;

[0062] y′ b-1 =y b-1 ⊙Silu(z b-1 )

[0063] ⊙ represents term-by-term multiplication, here we use z b-1 To control y b-1 The size of the change in each item;

[0064] X b = linear6(y′ b-1 )+X b-1

[0065] Finally, the skip connection is used to output the result X of this layer b , the intermediate variables X1, X3, ..., X 23 Enter the next stage network.

[0066] For f2(X b )=TRVim(X b ) are processed as follows:

[0067] Normalize the output of the previous layer:

[0068] X' b-1 =norm(X b-1 )

[0069] Then it passes through the Vim layer and is added to the part before the residual:

[0070] X” b-1 =X b-1 +Vim(X' b-1 )

[0071] Pass it through the normalization layer again:

[0072] X″′ b-1 =norm(X” b-1 )

[0073] Finally, after passing through the Gelu and MLP layers, the output is through the residual connection:

[0074] X b =X″ b-1 +linear mlp (Gelu(X″′ b-1 ))

[0075] For ScanningMamba, f3(X i )=SCVim(X i The specific processing of ) is as follows. Reverse the output of the previous layer and reverse the order:

[0076]

[0077] Enter both into the layer() function, which refers to the Mambamixer function described above:

[0078]

[0079] X d =layer(X d-1 )

[0080] Adjust the reverse output sequence to the forward direction and take the average output:

[0081]

[0082] The neck is the connection between the backbone network and the head. Its purpose is to further adjust the feature representation without changing the output size of the backbone network. It extracts higher-level semantic information by reducing the resolution or number of channels of the feature map, fuses multi-level information, and provides appropriate feature representation for the head. The neck is usually composed of convolutional structures such as convolutional layers and pooling layers. Figure 2 shown.

[0083] The intermediate features of the odd-numbered layer output of the lightweight Mamba are downsampled and connected with the convolution block through a residual connection. Finally, they are output through the full convolution layer and then through a simple feature pyramid network (FPN). Five feature layers of different sizes are output to the head part.

[0084] When the recommended Scan Mamba layer is selected, the matching neck network processing flow is as follows:

[0085] Y 2k-1 = Downconv(X 2k-1 )

[0086] Downconv() represents downsampling convolution, X 2k-1 It represents the output result of the 2k-1th layer, Y kk-1is the result of downsampling, k is an integer between 1 and 12;

[0087] Z 2k+1 =Y 2k-1 +Z 2k-1 +conv(Y 2k-1 +Z 2k-1 )

[0088] The above formula is the recursive representation of the front and back residual connections, Z 2k+1 It represents the intermediate features obtained after the k-th recursion, and finally the results of the 23rd layer are processed into a simple feature pyramid;

[0089] F feature1 = pooling(Z 23 )

[0090] F feature2 =Z 23

[0091] F feature3 =transpose(Z 23 )

[0092] F feature4 =transpose 2 (Z 22 )

[0093] F feature =[F feature1 ,F feature2 ,F feature3 ,F feature4 ]

[0094] F feature is the obtained multi-dimensional feature.

[0095] When the other two Mamba layers are selected, the matching neck network processing flow is as follows:

[0096] M 2p-1 = Downconv(X 2p-1 )

[0097] Downconv() represents downsampling convolution, X 2p-1 It represents the output result of the 2p-1 layer, M 2p-1 is the result of downsampling, p is an integer between 1 and 6;

[0098] N 2p+1 =M 2p-1 +N 2p-1 +conv(M 2p-1 +N 2p-1 )

[0099] The above formula is the recursive representation of the front and back residual connections, N 2P+1 It represents the intermediate features obtained after the Pth recursion, and finally the results of the 11th layer are processed into a simple feature pyramid;

[0100] F feature5 = pooling(N 11 )

[0101] F feature6 =N 11

[0102] F feature7 =transpose(N 11 )

[0103] F feature8 =transpose 2 (N 11 )

[0104] F feature_other =[F feature5 ,F feature6 ,F feature7 ,F feature8 ]

[0105] F feature_other is the obtained multi-dimensional feature.

[0106] The above operation is because the feature map of the last layer is too deep. It is the result of multiple extractions, and the surface information may be lost. Therefore, the prediction of the last layer (rich in global information) and the prediction of the shallower layer (rich in local information) are combined through the skip structure, and local prediction is performed while complying with the global prediction. If only the feature map of the last layer is passed through the full convolutional network to obtain a mask of the original image size, the final detection effect is relatively general, and it is difficult to converge. The residual layers before and after are also reminding the model to remember the previous and next associations.

[0107] The standard feature pyramid is divided into one column of upsampling and one column of downsampling. For the feature prediction before upsampling in each layer, the original size will be obtained after 2 times upsampling, and then it will be added to the part of the same size in the downsampling. Then the above operation is repeated again, and the sampling and addition are repeated, so that the multi-dimensional feature maps are fused with each other, retaining relatively more underlying information. The simple feature pyramid changes the order, first fusing the multi-dimensional features extracted in advance in some way, and then performing a series of up and down sampling at the same time to obtain a fixed number of feature maps of different sizes, which reduces the amount of calculation in the algorithm.

[0108] The head is built according to the structure of the large open source project mmdection. It is similar to the structure of the Mask-RCNN model.

[0109] Specifically, the head network of the present invention is divided into two types: segmentation head (seg-head) and decoder head (decoder-head), each of which is cascaded by the region proposal network head (Region Proposal Network Head, RPN Head) and the region of interest head (Region of interest Head, ROI Head). Here and below, except for the segmentation head (seg-head) and decoder head (decoder-head) designed by the present invention, all heads and loss functions are fixed functions under the large open source project mmdection.

[0110] The segmentation head (Seg-head) takes the RPN head with both input and output of 256 as the first step, and connects the cross entropy loss function (cls_loss) and the smooth L1 norm loss function (smoothl1loss) as the loss function of the detection box. The ROI Head is divided into two branches, one uses a single ROI function (singleROI) as an extractor (Extractor), uses the pooling layer function Alignpooling to pass the candidate features to the detection box head (bbox_head), and also uses the cross entropy loss function (cls_loss) and the smooth L1 norm loss function (smoothl1loss) as the loss function to fine-tune the result; the other is passed to the Fully Convolutional Networks Mask Head (FCN Mask Head) to predict the target mask.

[0111] The decoder head (decoder-head) also takes the RPNhead with both input and output of 256 as the first step, and connects the cross entropy loss function (cls_loss) and the smooth L1 norm loss function (smoothl1loss) as the loss function of the detection box. The ROI Head is divided into two branches, one uses a single ROI function (singleROI) as an extractor (Extractor), and uses the pooling layer function Alignpooling to pass the candidate features to the detection box head (bbox_head), but the other is not passed to the network mask terminal function (Fully Convolutional Networks Mask Head, FCN Mask Head), but is sent to a trainable prompt encoder (prompter encoder), and the results of both are sent to the image decoder (imagedecoder), unified output mask and use cross entropy to calculate the loss.

[0112] Step 3: Train the neural network. In the Pytorch deep learning framework, the backbone model is trained under unlabeled data and the complete network is trained under labeled remote sensing data. Figure 3 As shown in the figure, the specific training method of the former is: for the 50,000 augmented images, they are sent to the teacher and student models respectively, and the output results are all 256×64×64 matrices, flattened along the first dimension, and the norm loss is directly calculated for parameter return, and the final parameters are retained. When training the complete network, the training results of the backbone are retained as pre-training weights, the neck and head are connected, and all parameters are fine-tuned on the remote sensing data set. During the training process, the training rounds are set to 800 rounds (epochs). Each round can select a part of the training set and randomly arrange them. Every 3 rounds, the optimal network parameters are selected according to the test data. Finally, the optimal network parameters of the entire process are retained as the final training results.

[0113] Step 4: Instance segmentation of remote sensing images, using the lightweight remote sensing instance segmentation network model based on state space knowledge distillation trained in step 3 to perform instance segmentation in remote sensing scenarios. Since the present invention uses the large open source project mmdection as the code framework, after the user imports the trained parameters in step 3 into the lightweight remote sensing target instance segmentation network, directly run the official test file image_demo.py, and input the image path that needs instance segmentation into the model. The model will output the visualization result after the image instance segmentation at the same location, as well as a yaml type file storing the specific data of the mask and detection box. When you want to test mAP and other related indicators according to the standard data set matched during training, directly run the official test file test.py. Under this framework, the model will automatically process the images in the test set one by one and calculate the mAP indicator. When it is necessary to migrate to an application scenario that needs to process other types of targets, the training data set can be replaced with a data set containing the corresponding target category in step 3 according to the type of instance segmentation target required. When the model needs to be deployed on a physical object, the software can be configured according to the environment specified in the specific implementation mode of the present invention to run normally. According to the model design of the present invention, the model of the present invention ensures that the performance is similar to that of the basic model under the same task, while the required number of parameters is basically smaller than the common traditional model under the same task type, and is much smaller than the basic model under the same task. It can effectively solve the contradiction between limited onboard resources and small algorithm deployment space but high performance requirements.

[0114] Figure 1The structure of the backbone network is shown. From left to right below, three different backbone networks are designed. The number in the upper left corner indicates the number of stacked layers, which is 12 or 24 layers. The function boxes such as ViM and flip passed by the arrows indicate the functions or structures passed in the model calculation. The specific description is shown in the backbone network of step 2. The dotted rectangular box shows the overall structure of the model, but focuses on showing the position of the backbone network in the model. Among them, the gray lightweight Mamba encoder is the three different backbone networks. After the image is cut into small squares, it passes through the linear layer and adds position encoding, and passes through the backbone, neck and head in turn to get the final output.

[0115] Figure 2 The connection and construction of the neck and the front and back are shown. The red background is the downsampling layer mentioned above. Through the downconv downsampling function, the sampled result is sent to the blue background, that is, the residual connection stacking layer. The specific processing formula can be found in the neck description of step 2. Finally, the output result corresponding to the last layer to be sampled is sent to the yellow background, and after the simple feature pyramid layer in light blue, convolution and pooling are performed respectively, and the final multi-layer output can be sent to the head network.

[0116] Figure 3 The training method of the distillation method is shown. In chronological order, it is divided into SA-1B data augmentation, lightweight encoding head teacher-student model construction, and one-norm loss feedback. The pictures after data augmentation are sent to the teacher model above and the student model below. The teacher model uses the official parameters given by Facebook, and the final output of the two is solved for the L1 loss norm. Then the parameters of the teacher model are fixed, and the learning model parameters are adjusted by gradient descent according to the result of the loss function, and the final training results are retained, as shown in step 3.

[0117] Figure 4a-4f The actual effect of using this algorithm is demonstrated. Figure 4a and 4b The test results on the WHU dataset are shown. Figure 4c and Figure 4d The test results on the NWPU dataset are shown. Figure 4e and 4f The test results of SSDD data integration are shown. In general, whether it is for dense or large buildings ( Figure 4a and Figure 4b ) or a scene with mixed multiple categories of targets ( Figure 4c and Figure 4d ) or a docked ship or discrete ship target with a complex background ( Figure 4e and Figure 4f ), the models of the present invention all have good effects.

Claims

1. A method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation, characterized by: The steps include: Step 1: Computer reading data: Use a computer to read the unlabeled SA-1B data set and the labeled remote sensing image data, and pre-process the image data respectively; Step 2: Construct a knowledge distillation network based on the state space model and a lightweight remote sensing target instance segmentation network; The state-space model-based knowledge distillation network consists of two parts: the teacher network and the student network. The lightweight remote sensing target instance segmentation network consists of three parts: the backbone, the neck, and the head. Among them, the teacher network is the image encoder of the SAM basic model, and no parameter adjustment is performed during training; the backbone network is the student network obtained after training based on the state-space model knowledge distillation network, which is responsible for extracting multi-level image information; the neck network is used for multi-level feature fusion; and the head is responsible for processing specific tasks. Step 3: Train the neural network; After constructing the knowledge distillation network based on the state-space model using the Pytorch deep learning framework, use the unlabeled data to train this network, save the training parameter results, import the student network based on the state-space model knowledge distillation network, and use it as the backbone network of the lightweight remote sensing target instance segmentation network. After connecting the neck and head, use the labeled remote sensing data to train the lightweight remote sensing target instance segmentation network; The specific training method is: in the first training, construct the backbone through the state-space model, and then use it as a student model to train the lightweight method of knowledge distillation to make it learn the complex knowledge of the teacher model, and then retain the training results of the backbone as the pre-training weights, connect the neck and head, and fine-tune all parameters on the remote sensing data set; During the training process, the optimal network parameters for lightweight remote sensing target instance segmentation are selected according to the test data, and the network parameters for lightweight remote sensing target instance segmentation are saved as the final training results. Step 4: Remote sensing image instance segmentation: Use the lightweight remote sensing target instance segmentation network model trained in step 3 to achieve high-quality remote sensing image instance segmentation on the test set or actual image.

2. According to claim 1, a method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation is characterized by: In step 1, the data sources used are the official SAM SA-1B dataset and the well-known public remote sensing datasets SSDD, WHU, and NWPU datasets; in the SA1B dataset, 12,000 images are selected, data augmentation is performed on them, and they are expanded to 50,000 1024x1024 images; the SSDD, WHU, and NWPU datasets are all small and medium-sized datasets, and the three are respectively for optical SAR ships, static building targets, and mixed scenes of multiple types of targets; the complete dataset is selected and uniformly converted to the COCO format to adapt to the MMCV framework, and data augmentation is also performed, and it is expanded to a size of 1024x1024 without changing the aspect ratio.

3. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 1, characterized in that: In step 2, three different backbone networks are designed. The VanillaMamba consists of 12 layers of normal visual Mamba blocks Vim Block. The model is the smallest but has weak learning ability. The TransMamba consists of 12 layers of changed visual Mamba blocks TRVim Block combined with Transfromer, which specifically includes five components: Residual Layer, Normalization Layer, Mamba Layer, Gelu activation function layer and Multilayer Perceptron MLP. It has strong learning ability, but the model is the largest and the reasoning is slow. The ScanningMamba consists of 24 layers of SCVim Block. The model size is medium, but it has strong learning ability and fast reasoning. Users can choose according to task requirements.

4. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 3 is characterized in that: All three backbone networks divide each input image into I = n 2 The p×p×c slices are flattened and multiplied by the training matrix W and the sine and cosine position codes to obtain the matrix If X is specified l refers to the output of layer I, then for VanillaMamba, X b =X b-1 +Vim(X b-1 ), b = 1, 2, ..., 12; for TransMamba, X b =X b-1 +TRVim(X b-1 ), b = 1, 2, ..., 12; for ScanningMamba, X d =X d-1 +SCVim(X d-1 ), d = 1, 2, ..., 24; Among them, Vim() is the original visual state space model stacking layer, TRVim(X i ) During processing, the output of the previous layer is normalized, X′ b-1 =norm(X b-1 ); then pass through the Vim layer and add it to the part before the residual to get X″ b-1 =X b-1 +Vim(X′ b-1 ), and then pass through the normalization layer again: X″′ b-1 =norm(X″) b-1 ); finally, after passing through the Gelu and MLP layers, the residual connection is output to obtain X″ b-1 +linear mlp (Gelu(X″′ b-1 )); while SCVim(X i ) When processing, first reverse the output of the previous layer and reverse the order to get And input the two into the layer constructed by the official mambamixer function, and get and X d =layer(X d-1 ), and finally adjust the reverse output sequence to the forward direction And take the average and output 5. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 4, characterized in that: Neck is the connection structure between the backbone network and the head. When selecting the recommended scanning mamba layer, use an odd number of layers X 2k-1 The intermediate features of the output, X 2k-1 It represents the output result of the 2k-1th layer, and we get Y 2k-1 = Downconv(X 2k-1 ), Y 2k-1 It is the result of downsampling, k is an integer between 1 and 12, and then it is output through the full convolution layer, specifically expressed as Z 2k+1 =Y 2k-1 +Z 2k-1 +conv(Y 2k-1 +Z 2k-1 ), Z 2k+1 represents the intermediate features obtained after the kth recursion. Finally, the results of the 23rd layer are processed into a simple feature pyramid to obtain F feature =[pooling(Z),Z,transpose(Z),transpose 2 (Z)], where Z represents the output of the last odd layer, pooling represents pooling, and transpose represents deconvolution.

6. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 5, characterized in that: When choosing the other two Mamba layers, use the odd layer X 2p-1 The intermediate features of the output, X 2p-1 It represents the output result of the 2p-1 layer, and M 2p-1 = Downconv(X 2k-1 ), M 2p-1 It is the result of downsampling, p is an integer between 1 and 6, and then it is output through the full convolution layer, specifically expressed as N 2p+1 =M 2p-1 +N 2p-1 +conv(M 2p-1 +N 2p-1 ), N 2P+1 It represents the intermediate features obtained after the Pth recursion. Finally, the results of the 11th layer are processed into a simple feature pyramid to obtain F feature =[pooling(Z),Z,transpose(Z),transpose 2 (Z)], where Z represents the output of the last odd layer, pooling represents pooling, and transpose represents deconvolution.

7. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 1, characterized in that: In step three, under the Pytorch deep learning framework, during the first training, the backbone is constructed through the state space model, and it is used as the student model for training using the lightweight method of knowledge distillation. The 50,000 augmented images are sent to the teacher and student models respectively, and the output results are both matrices of 4, which are flattened along the first dimension, and the one-norm loss is directly calculated for parameter feedback to retain the final parameters.

8. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 1 or 7, characterized in that: In step three, under the Pytorch deep learning framework, during the second training, i.e., fine-tuning, the training results of the backbone are retained as pre-trained weights, the neck and head are connected, and all parameters are fine-tuned on the remote sensing dataset. During the training process, the training rounds are set to 800 epochs. Each epoch selects a part of the training set and arranges them randomly. Every 3 rounds, the optimal network parameters are selected according to the test data. Finally, the optimal network parameters of the entire process are retained as the final training results.

9. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 8, characterized in that: During training, the AmpOptimWrapper optimization wrapper type of mmdection is used, and bfloat16 floating-point data is used for operation; the AdamW optimizer is used, and the learning rate base_lr and weight decay rate are set to 0.05; the linear learning rate scheduler is LinearLR and the cosine annealing learning rate scheduler is CosineAnnealingLR. The former starts from iteration 0 and continues to iteration 50, with an initial learning rate factor of 0.001, and the scheduling is based on the number of iterations rather than the number of epochs. The latter starts from the first epoch and continues until the maximum number of rounds max_epochs ends; it The minimum learning rate eta_min is base_lr*0.001, and the cosine annealing period T_max is also set to max_epochs. The scheduling is based on the number of epochs. The learning rate is dynamically adjusted according to the number of iterations or epochs to optimize the training process of the model. In addition, the image and the corresponding annotations are randomly flipped horizontally with a probability of 50%, and the image size is randomly adjusted using the RandomResize function. The scale varies based on the crop_size, and the aspect ratio varies within (0.1, 2.0). The filter FilterAnnotations is responsible for filtering out annotation information whose bounding box or mask size is less than the given threshold 1e-5.

10. The method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation according to claim 1, characterized in that: In step 4, the instance segmentation in remote sensing scenarios is performed using the lightweight remote sensing instance segmentation network model based on state-space knowledge distillation trained in step 3. After the user imports the parameters trained in step 3 into the lightweight remote sensing target instance segmentation network, directly run the official test file image_demo.py, and input the image path that needs instance segmentation into the model. The model will use the lightweight state-space model backbone to obtain features, and then use Neck feature fusion, and use Mask-RCNN-Head to predict the final result. The model will output the visualization result after image instance segmentation at the same location, as well as a yaml file that stores the specific data of the mask and detection box. Test on the SSDD, WHU, and NWPU datasets. The target size scale of each test set image includes ships and aircraft targets. When you want to test mAP-related indicators based on the standard dataset matched during training, directly run the official test file test.py. Under this framework, the model will automatically process the images in the test set one by one and calculate the mAP indicator.

Images