A method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation
Through the knowledge distillation method based on the state space model, a lightweight remote sensing target instance segmentation model is constructed, which solves the problems of slow inference speed and large number of parameters in remote sensing target detection, and realizes efficient remote sensing target detection under limited satellite resources.
Patent Information
- Application Number
- CN202411183072.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-08-27
- Publication Date
- 2025-10-24
- Estimated Expiration
- 2044-08-27
AI Technical Summary
In remote sensing target detection, existing deep learning models have the disadvantages of slow inference speed and large number of parameters, making it difficult to achieve a balance between real-time performance and accuracy under the condition of limited satellite resources. Traditional lightweight methods cannot effectively improve model performance.
A knowledge distillation method based on the state-space model is used to construct a lightweight remote sensing target instance segmentation model. Through pre-training and fine-tuning strategies, the knowledge of the large model is distilled into a small student model. Combined with the expandability of the state-space network layer, multi-layer stacking results are quickly calculated, reducing model storage and inference time.
It achieves faster inference speed and smaller storage space in remote sensing missions while maintaining detection results similar to conventional models, and is suitable for environments with limited onboard resources.
Smart Images

Figure CN119274051B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The application relates to a construction method of a light remote sensing target instance segmentation model based on state space model knowledge distillation, and belongs to the field of satellite networks. BACKGROUND
[0002] A state space model (SSM) is a model that describes the state before the model and predicts the state at the next moment according to the current input. The SSM is first applied to the control field, taking a continuous sequence as input and predicting an output sequence. The model has two equations, a state equation and an output equation. The state equation is used to describe the multiplication of a matrix and an input, and then the result of the multiplication of a matrix and a previous state. The output equation describes how the state is converted into the output through a matrix, and how the input affects the output through a matrix.
[0003] Albert Gu et al. designed the S4 model to solve the quadratic complexity of the attention mechanism in the Transformer architecture and the sequence forgetting problem of the RNN at the same time, so as to obtain a model with the advantages of both. The S4 model discretizes the original SSM model, that is, uses a discrete sequence as input. The S4 model introduces a time step and a HiPPO memory mechanism. When a discrete signal is received each time, the signal value is retained until a new discrete signal is received. The processing interval of the operation is a new learnable parameter, called time step. In this way, the model is kept in stages. According to the learned time step, the input signal is sampled, and the final output can be generated. In the discretization operation, there are two ways of zero-order holding 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 classical image recognition task, the convolution kernel is used to aggregate the features, and through expression expansion, it can be written in a form similar to convolution. Therefore, it can be trained in a convolution mode during training and inferred in a loop mode during testing, greatly improving 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 seen so far into a coefficient vector, which can well capture the recent tokens and decay the old token state representation. Through function approximation, the optimal solution of the state matrix is generated.
[0005] Albert Gu et al. proposed Mamba model on the basis of the previous work, which is a collection of the previous work. Mamba is inspired by the Simplified State Space Layers for Sequence Modeling (S5), which uses a hardware-aware algorithm to reduce IO access between computer levels. It also introduces selective information processing of structured SSM (S6), which is represented by a learnable input mapping function at the bottom level, which produces different calculation parameters for different inputs. Mamba also borrows the concept of LSTM and gate circuit, combining SSM structure and Transformer gate MLP block into a layer to simplify the deep sequence model architecture of the past. In terms of task performance, Mamba has achieved better results than SOTA in language, audio, and DNA sequence modalities, and has surpassed the Transformer of the same size in language tasks, and is close to the Transformer twice as large.
[0006] Mamba model has attracted widespread attention since its introduction, and many scholars have extended it to other directions, leading to the development of Vision Mamba (ViM). Similar to Vision Transformer (ViT), ViM first divides the input image into small patches, then flattens the small patches into a sequence and maps them into tokens, and then applies position encoding and adds a class tensor to unify the input of the stacked layers of ViM. In each layer, two different linear mappings are performed before and after, and the sequence is processed in both forward and backward directions in between. Finally, the Sigmoid linear unit (SiLU) activation function is used, and it is worth noting that for non-classification tasks, no additional learnable classification token is needed. Finally, ViM has better results and efficiency in benchmark tests such as ImageNet classification, COCO object detection, and ADE20K semantic segmentation. For high-resolution images, ViM reduces GPU memory by 86% compared to DEIT while speeding up by 2.8 times.
[0007] In the past few months, Mamba has been roughly divided into the first and second stages 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, implemented the fusion of S6 and gated multilayer perceptron (MLP), and provided an open-source network level for use. Inspired by this, Vision Mamba and VMamba appeared three months later, both of which improved it from different angles and became the main two types of improvement ideas for subsequent researchers. The former, proposed by L. Zhu et al. in January 2024, combined the idea of Transformer, replaced the original Mamba layer with a stack of Transformer layers, and used image slicing and encoding to process, achieving a visually effective Vision Mamba. The latter, proposed by Y. Liu et al. at almost the same time, mainly considered that Mamba lacks the viewing order of image slices, and the original slicing method may cause images that are close together to be divided far apart. Therefore, they proposed the idea of fast scanning, using SS2D Block for multiple scanning methods, and cooperating with the internal VSSBlock for stacking operations. It is worth noting 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 SS2D Block that puts the scanning order at the front. Subsequently, Mamba Mixer was proposed for visual Mamba, and Mamba2 was proposed for traditional Mamba, establishing a mathematical connection between Mamba and Transformer.
[0008] Model lightening (Distributed computing) refers to the method of reducing model storage and speeding up inference to obtain a more efficient model within the allowed range without affecting the performance of the model. In many practical application scenarios, it can play a good role in practice. For example, in the satellite remote sensing scene, remote sensing satellites have been applied to many scenarios, and new low-orbit giant constellation networks have made up for the shortcomings of traditional high-orbit satellite resolution or real-time detection difficulty, and have become a research hotspot for various countries. However, there is a contradiction between real-time performance and speed accuracy in remote sensing target detection, especially in low-orbit satellite detection. Satellite space resources are limited, and complex deep learning models are difficult to deploy, while simple models are difficult to cope with complex remote sensing scenarios, so it is crucial to design a lightweight target detection algorithm suitable for remote sensing scenarios.
[0009] Then, looking at the target detection field in 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 visual base model, also known as the large model, has made excellent performance and good scalability and generalization in various visual tasks through a substantial expansion of data and model size, and has become a major research hotspot at present. But the performance of the visual large model is good, the inference speed is slow, and the parameter quantity is too large, for example, the target segmentation SAM large model uses more than 1 billion parameters. The goal is to use lightweight methods to deploy the base model on the star under the premise of ensuring accuracy, substantially compressing parameters and power consumption.
[0010] The model lightweight method specifically includes three kinds: knowledge distillation that transfers knowledge from a large teacher model to a small student model, model pruning that takes network sparsification as the core idea, and model quantization that converts floating-point algorithms to fixed-point from the storage perspective. But pruning cannot fundamentally speed up model inference, and quantization generally cannot achieve a magnitude of improvement, so research is started from the perspective of knowledge distillation.
[0011] The present application proposes a lightweight remote sensing target instance segmentation model construction method based on state space model knowledge distillation, which constructs a backbone network with strong learning ability according to the state space model, uses knowledge distillation and learns the complex knowledge of a large network using ViT as a teacher model, quickly calculates the result of multi-layer stacking through the expandability of the state space network layer to speed up the inference speed, and the student network constructed itself substantially reduces the storage, and the neck network theoretically alleviates the loss of small targets to some extent. The time of the total process of model inference and the space occupied by model storage are greatly reduced, the time cost and space cost required for model running can be greatly reduced, and a practical solution is proposed for the contradiction between limited spaceborne resources and high model performance requirements. The present application constructs a lightweight construction method of a remote sensing instance segmentation model based on state space model knowledge distillation for space remote sensing problems, which can segment the mask while detecting the target, and provides a referenceable configuration method. In addition, this technology can also be easily migrated to other task scenarios that need to be lightweight, such as vehicle detection, agricultural monitoring, and aerial ground detection. SUMMARY
[0012] The purpose of the present application is to design a lightweight remote sensing target instance segmentation model construction method based on state space model knowledge distillation. Using this method, the network model can have faster speed and smaller storage than the conventional model without distillation, and can be efficiently trained, tested and used on a conventional graphics processing unit (GPU), and can achieve the same effect as the conventional model. Specifically, the present application is implemented by the following technical solutions:
[0013] The present application is a lightweight remote sensing target instance segmentation model construction method based on state space model knowledge distillation, which essentially belongs to the strategy of "pre-training + fine-tuning". This method constructs lightweight coding backbone to learn the knowledge of large model as pre-training weight; then connects the neck and detection head, and sets all parameters as trainable, and then fine-tunes on the remote sensing dataset to achieve the expected 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 backbone network is still not frozen, which ensures its specialization in 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 SA-1B dataset without annotation and the remote sensing image data with annotation, and pre-process the image data respectively.
[0015] Step 2: Construct state space model knowledge distillation network and lightweight remote sensing target instance segmentation network.
[0016] The present application proposes a construction method of state space model knowledge distillation network and lightweight remote sensing target instance segmentation network. The state space model knowledge distillation network is composed of two parts: teacher network and student network, and the lightweight remote sensing target instance segmentation network is composed of three parts: backbone, neck and head. Among them, the teacher network is the image encoder of the SAM base model, using the parameters provided by the official, and no parameter adjustment is made during training; the backbone network is the student network obtained after training the state space model knowledge distillation network, responsible for extracting multi-level image information, and essentially a high-dimensional feature extractor; the neck network is used for multi-level feature fusion; the head is responsible for processing specific tasks.
[0017] Step three: training the neural network. After constructing the knowledge distillation network based on the state space model using the Pytorch deep learning framework, the network is trained using unlabeled data, and the trained parameter results are saved. After importing the student network based on the state space model knowledge distillation network, it is used as the backbone network of the lightweight remote sensing target instance segmentation network, and after connecting the neck and head, the labeled remote sensing data is used to train the lightweight remote sensing target instance segmentation network (i.e. the backbone, neck and head network). The specific training method is as follows: in the first training, the backbone is constructed by the state space model, and then it is used as a student model to train the lightweight method based on knowledge distillation, so as to learn the complex knowledge of the large teacher model, and then the training results of the backbone are reserved as pre-training weights, the neck and head are connected, and all parameters are fine-tuned on the remote sensing data set to ensure its specialization on the remote sensing task, so as to finally achieve the expected effect. In the process of training, 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 four: remote sensing image instance segmentation. The lightweight remote sensing target instance segmentation network model trained in step three can realize high-quality remote sensing image instance segmentation on the test set or actual image.
[0019] The advantages of the present application are as follows: the present application uses the knowledge distillation method 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 in training avoids the possibility that the student model will produce large deviations after feature fusion due to slight differences with the teacher model during learning, and at the same time, the required data set can be selected for special training according to the specific task and object, which ensures the adaptability of the model under specific tasks and increases the value in practical application. BRIEF DESCRIPTION OF DRAWINGS
[0020] Figure 1 is the network overall structure diagram of the method of the present application.
[0021] Figure 2 is the neck network structure diagram of the method of the present application.
[0022] Figure 3 is the distillation training structure diagram of the method of the present application.
[0023] Figures 4a-4f is the image instance segmentation diagram of the method of the present application. DETAILED DESCRIPTION
[0024] In order to better understand the technical solutions of the present application, the following will combine the drawings of the present application with the specific embodiments of the present application to make a further description of the present application. Figures 1-4f The specific embodiments of the present application are further described as follows:
[0025] The application is implemented by using Python programming under the Pytorch framework. First, the training and test data are prepared, then the network is built and the related parameters are configured; then the network is trained using the training data to obtain the optimal network parameters; finally, the test data is tested using the trained network model parameters to obtain the building extraction result. In 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 Python 3.10, Pytorch 2.1.2, Mmcv 2.1.0, Mamba-ssm 1.1.1, and causal-conv1d 1.1.1.
[0026] Step one: the computer reads data. The data used by the application is from the SA-1B dataset of SAM official and the famous public remote sensing dataset SSDD, WHU and NWPU dataset. The SA-1B dataset includes more than 1.1 billion segmentation masks collected on about 11 million licensed and privacy protected images. About 1.2w pictures can basically meet the requirements, which are basically data enhanced and expanded to about 5w 1024x1024 images. The SSDD, WHU and NWPU datasets are small and medium-sized datasets, and the three are respectively for optical SAR ship, building static target and multi-type target mixed scene. Select the complete dataset and convert it to COCO format to adapt to the framework of Mmcv, also perform data augmentation, and expand to 1024x1024 size without changing the aspect ratio.
[0027] Step two: construct a knowledge distillation network based on a 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 quantization of floating-point number fixed-point storage which are based on the core idea of sparsification, distillation can extract knowledge from large but effective models and teach it to small student models. After distillation, the new model can have the complexity of the original small model and the complex knowledge of the large model. In recent years, as a new method introduced into computer vision, the space state model (SSM) has produced impressive performance in current perception, generation and other tasks, and its unique expandability when stacked has significantly improved the processing speed of feature extraction. While the large model derived from the traditional vision transformer has made outstanding performance in various visual tasks through a substantial expansion of data and model size, and has played a promoting role in the current direction. However, although the performance of the visual large model is good, it is slow in reasoning and has a huge number of parameters, for example, the SAM large model for target segmentation uses more than 1 billion parameters. The ultimate goal is to use lightweight methods to make small models have the performance of basic models under the premise of ensuring accuracy, greatly compressing parameters and power consumption, and finally deploy them in various places.
[0029] The application proposes a construction method of a state space model knowledge distillation network and a lightweight remote sensing target instance segmentation network. The state space model knowledge distillation network is composed of a teacher network and a student network, and the lightweight remote sensing target instance segmentation network is composed of a backbone, a neck and a head. The teacher network is an image encoder of the SAM basic model, and the official parameters are used without parameter adjustment during training; the backbone network is the student network obtained after training of the state space model knowledge distillation network, 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 application, responsible for extracting features from the original input, as shown in Figure 1 It is composed 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 picture, and process the image information layer by layer. Similar to the foregoing, the output of the backbone network is a four-dimensional feature map. For each batch, for each picture in the batch, it is expanded to h×h×c size by clip(), and is cut into multiple p×p×c slices, 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 a 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 consists of 12 layers of normal Vision Mamba Blocks (Vim Block), the smallest model, but weaker learning ability; Trans Mamba consists of 12 layers of Transformer Vision Mamba Block (TRVim Block) combined with Transformer, similar to the structure of Transformer, specifically including five components: Residual Layer, Normalization Layer, Mamba Layer, Gelu activation function layer and Multilayer Perceptrons (MLP). It has strong learning ability, but the model is the largest and the reasoning is slow; Scanning Mamba consists 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 the task requirements, but the present invention recommends Scanning Mamba. The present invention is described in detail below from two aspects: intuitive and mathematical details.
[0032] The VanillaMamba was inspired by the original Mamba. It simply consists of 12 stacked layers of normal Vision Mamba Blocks (Vim Blocks). However, during the training of the present invention, although it can converge in the first stage, its model is relatively simple. At the same time, the present invention does not change the parameters related to the original memory. The graphics it processes are expanded from the original 256×256 size to 1024×1024, and the number of sequences has increased by 16 times. Therefore, its learning ability is relatively weak.
[0033] TransMamba greatly increases the parameters compared to the original backbone network, but the efficiency is also reduced to some extent. In addition, in the code implementation, the present application uses the SAM image encoder given in the third-party library MmCV and directly changes its stacking layer, resulting in a larger Python file when storing the model, because some parameters that the original Transformer needs but the present application does not need are also saved. When doing simple tests, the basic performance of this model is that the bounding box (BBox) can barely converge, but the segmentation part performs poorly. In addition, the present application does not replace the complete stacking layer, but only replaces the attention layer in the designed Transformer Vision Mamba Block (TRVim Block), modifies the activation function, and modifies the loss function to L1 norm to achieve better results, because taking the square may cause the value to be too small, which is not conducive to parameter backpropagation training.
[0034] ScanningMamba focuses on the scanning method in the model design. First, the stacking layer is changed to the Mamba mixer provided by Facebook instead of the original SAM image encoder, and the number of stacking layers is increased to 24 to enhance the learning ability. Second, the size of the embedding dimension is adjusted to reduce the dimension of the word vector and reduce the parameters of the model. Finally, the scanning strategy of RSMamba is selected, but the random scanning layer is not used, but only the forward and backward scanning methods are used. Specifically, when a series of word vectors are input, the order of the forward and backward sequences is reversed, and the parameters are shared when inputting each layer. After output, the reversed sequence is reversed again and added to the average of the forward sequence, so that the output of the layer is the same as the dimension of the input, and then it is sent to the next layer.
[0035] For the J slices obtained, use resize() to flatten them into a single dimension but a vector of length p 2 ×c, respectively denoted as where i represents the i-th picture in the batch, 1≤i≤batch, k represents the vector corresponding to the k-th slice, and the arrangement order is to scan each slice vector row by row, multiply each slice vector by the trainable matrix W, W is a learnable matrix with column number p 2 c, and row number d, and add the sine position encoding to each vector to obtain matrix X0
[0036]
[0037] wherein X position Learnable position encoding.
[0038] Note that the original formula of Vim is i.e. the column vector needs to be combined with the learnable class vector x class Concatenation, but here x class is not added. This is because the task of the present application has been distinguished from the original classification task, and secondly in the down-sampling stage of the neck, only even-numbered slices can use two-dimensional convolution down-sampling.
[0039] After that, X0 is input into the stacked layer, i.e. 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] wherein Vim() is the normal vision Mamba block mentioned in the foregoing, TRVim() is the vision Mamba block combined with Transfromer mentioned in the foregoing (Transformer Vision Mamba Block, TRVim Block), and SCVim() is the scanning Mamba block mentioned in the foregoing (Scanning Vision Mamba Block, SCVim Block), X b refers to the output result of the bth layer, b is a positive integer between 1 and 12, X d refers to the output result of the dth layer, d is a positive integer between 1 and 24.
[0046] For f1(X i ) = Vim(X iIn terms of the specific processing of the above, the following is the specific processing of the above:
[0047] The output of the previous layer is normalized:
[0048] X' b-1 = norm(X b-1 )
[0049] Where X b-1 represents the output of the b-1 layer, norm() represents the normalization function, and X' b-1 is the normalized result;
[0050] The normalized result is mapped to variables t b-1 and z b-1 by the linear() linear mapping function;
[0051] t b-1 = linear1(X' b-1 )
[0052] z b-1 = linear2(X' b-1 )
[0053] Where t b-1 is used to participate in normal forward and reverse operations, and z b-1 is responsible for how the results of forward and reverse are fused;
[0054] The parameter matrices B b-1 , C b-1 and Δ b-1 are then constructed in the following way;
[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] Where Silu() is the Sigmold activation function, linear p () is a different linear mapping function, p is a positive integer between 1 and 5, and the remaining symbols are as explained above;
[0060]
[0061] With is a discretized matrix, g() is a discretization function, X b is the output of the b-th layer, the intermediate variables X1, X3,..., X 23 is input to the next layer, SSM() is the state space model formula constructed with it, and t' is substituted into the above equation to obtain b-1 participate in the calculation, and finally the output y is obtained b-1 ;
[0062] y' b-1 = y b-1 ⊙Silu(z b-1 )
[0063] ⊙ denotes term-by-term multiplication, and z b-1 is used to control the change size of each term of y b-1 ;
[0064] X b = linear6(y' b-1 ) + X b-1
[0065] Finally, the result X b of this layer is output using the skip connection, and the intermediate variables X1, X3,..., X 23 are input to the next stage network.
[0066] The specific processing of f2(X b ) = TRVim(X b ) is as follows:
[0067] The output of the previous layer is normalized:
[0068] X' b-1 = norm(X b-1 )
[0069] Then pass through the Vim layer and add it to the part before the residual error:
[0070] X" b-1 = X b-1 + Vim(X' b-1 )
[0071] Again, pass through the normalization layer:
[0072] X""" b-1 = norm(X" b-1 )
[0073] Finally, pass through the Gelu and MLP layers and output through the residual link:
[0074] X b = X" b-1 + linear mlp (Gelu(X" b-1 ))
[0075] For ScanningMamba, f3(X i ) = SCVim(X i ) is processed as follows. The output of the previous layer is reversed, and the order of the front and back is reversed:
[0076]
[0077] Both are input into the layer() function, which refers to the Mambamixer function described above:
[0078]
[0079] X d = layer(X d-1 )
[0080] The reversed output sequence is adjusted to be forward, and the average is taken as the output:
[0081]
[0082] The neck is the connecting part between the trunk network and the head, and its purpose is to further adjust the feature representation without changing the output size of the trunk network, extract higher-level semantic information by reducing the resolution or channel number of the feature map, fuse multi-level information, and provide appropriate feature representation for the head. The neck is usually composed of convolutional structures such as convolutional layers, pooling layers, etc., as shown in the accompanying drawings. Figure 2
[0083] Using the intermediate features of the odd layer output of the lightweight Mamba, it is down-sampled and connected with the residual connection after the convolution block, and finally output through the full convolution layer, and then through the simple feature pyramid (Feature Pyramid Network, FPN), 5 feature layers of different sizes are output to the head (Head) part.
[0084] When the recommended scanning Mamba layer is selected, the matched neck network processing flow is as follows:
[0085] Y 2k-1 = Downconv(X 2k-1 )
[0086] Downconv() represents down-sampling convolution, X 2k-1 represents the output result of the 2k-1 layer, and Y kk-1 is the result of down-sampling, 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 rear residual connection, Z 2k+1 represents the intermediate feature obtained after the kth recursion, and finally the result of the 23rd layer is simply a 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 matched neck network processing flow is as follows:
[0096] M 2p-1 = Downconv(X 2p-1 )
[0097] Downconv() represents down-sampling convolution, X 2p-1 represents the output result of the 2p-1th layer, M 2p-1 is the result of down-sampling, and 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 rear residual connection, N 2P+1 represents the intermediate feature obtained after the Pth recursion, and the result of the 11th layer is finally simply 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, which is the result of multiple extractions, and the surface information may be lost. Therefore, by combining the prediction of the last layer (rich in global information) and the prediction of the shallower layer (rich in local information) through the skip structure, local prediction is performed while complying with global prediction. If only the feature map of the last layer is obtained after passing through the full convolutional network to obtain the original size of the mask, the final detection effect is general, and there is difficulty in convergence. The residual layers before and after remind the model to remember the correlation before and after.
[0107] The standard feature pyramid is divided into one column of up-sampling and one column of down-sampling. For the feature prediction before each layer of up-sampling, after 2 times of up-sampling, the original size is obtained, and then it is added to the part with the same size in the down-sampling. Then the above operation is repeated, and repeated sampling and addition are performed to make the multi-dimensional feature maps fuse with each other, and relatively more bottom layer information is retained. The simple feature pyramid changes the order, first fuses the multi-dimensional features extracted in advance in a certain way, and then performs a series of up-sampling and down-sampling to obtain a fixed number of feature maps with different sizes, which reduces the computational amount in the algorithm.
[0108] The head (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 (Mask-RCNN).
[0109] Specifically, the head network of the present application is divided into two types: segmentation head (seg-head) and decoder head (decoder-head), each of which is formed by cascading a Region Proposal Network Head (RPN Head) and a Region of Interest Head (ROI Head) in front and back. Here and hereinafter, all the heads and loss functions are fixed functions under the large open source project mmdection, except for the segmentation head (seg-head) and decoder head (decoder-head) designed by the present application.
[0110] The segmentation head (Seg-head) takes the RPN head with an input and output of 256 as the first step, and outputs the cross-entropy loss function (cls_loss) and the smooth L1 norm loss function (smoothl1loss) as the loss function of the detection frame. The ROI Head is divided into two branches, one of which takes a single ROI function (singleROI) as an extractor (Extractor) and uses a pooling layer function Alignpooling to pass the candidate features to the detection frame 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 results; the other branch is passed to the Fully Convolutional Networks Mask Head (FCN Mask Head) for predicting the target mask.
[0111] The decoder head (decoder-head) also takes the RPN head with an input and output of 256 as the first step, and outputs the cross-entropy loss function (cls_loss) and the smooth L1 norm loss function (smoothl1loss) as the loss function of the detection frame. The ROI Head is divided into two branches, one of which takes a single ROI function (singleROI) as an extractor (Extractor) and uses a pooling layer function Alignpooling to pass the candidate features to the detection frame head (bbox_head), but the other branch is not passed to the Fully Convolutional Networks Mask Head (FCN Mask Head), but is sent to the trainable prompter encoder, and the results of the two are sent to the image decoder (image decoder), and the mask is output uniformly and the cross-entropy loss is used.
[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 for the former is as follows: 50,000 augmented images are fed into the teacher and student models respectively, with the output results being 256×64×64 matrices. These are flattened along the first dimension, and the parameters are directly passed back using a norm loss, retaining the final parameters. When training the full network, the backbone training results are retained as pre-trained weights, the neck and head are connected, and all parameters are fine-tuned on a remote sensing dataset. During training, the training rounds are set to 800 rounds (epochs). Each round can select a portion of the training set and randomly arrange them. Every three rounds, the optimal network parameters are selected based on the test data. Finally, the optimal network parameters from the entire process are retained as the final training result.
[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 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 output the visualization result after the image instance segmentation at the same location, as well as a yaml type file that stores the specific data of the mask and detection box. When you want to test relevant indicators such as mAP based on 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 of the present invention to ensure normal operation. 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 that of 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 on-board resources and small algorithm deployment space but high performance requirements.
[0114] Figure 1The structure of the backbone network is shown, and the three different designed backbone networks are shown from left to right below, and the number in the upper left corner represents that the number of stacked layers is 12 or 24, and the arrow passes through the ViM, flip and other function boxes, which represents the function or structure passed in the model calculation, and the specific description is shown in the backbone network in step two. The overall structure of the model is shown in the dashed rectangular box, but the position of the backbone network in the model is focused on, wherein the gray lightweight mamba encoder is the three different backbone networks. After the image is divided into small squares, the position coding is added through the linear layer, and then the backbone, neck and head are sequentially passed to obtain the final output.
[0115] Figure 2 The connection and construction method of the neck and the front and rear are shown, and the red background is the down-sampling layer mentioned in the foregoing, and the sampled result is sent to the blue background, that is, the residual connection stacking layer, through the downconv down-sampling function, and the specific processing formula is shown in the description of the neck in step two. Finally, the output result corresponding to the last layer of the sampled is sent to the yellow background, and the light blue simple feature pyramid layer is used for convolution and pooling, so that the final multi-layer output can be sent to the head network.
[0116] Figure 3 The training method of the distillation method is shown, and in time sequence, it is divided into SA-1B data augmentation, lightweight coding head teacher-student model construction, and one-norm loss back propagation. The pictures after data augmentation are sent into the upper teacher model and the lower student model, the teacher model adopts the official parameters given by Facebook, and the final output of the two is solved L1 loss norm. Then, the parameters of the teacher model are fixed, and the learning model parameters are adjusted according to the loss function result through gradient descent, and the final training result is reserved, as shown in step three.
[0117] Figures 4a-4f The actual effect of using the algorithm is shown. Figure 4a and 4b The test results on the WHU data set are shown, Figure 4c and Figure 4d The test results on the NWPU data set are shown, Figure 4e and 4f The test results on the SSDD data set are shown. In summary, whether it is for dense or large buildings Figure 4a and Figure 4b ), or a mixed scene of multiple target categories Figure 4c and Figure 4d ), or a complex background of a moored ship or a discrete ship target Figure 4e and Figure 4f ), the model of the application has good effect.
Claims
1. A method for constructing a lightweight remote sensing target instance segmentation model based on state space model knowledge distillation, characterized in that: Comprising the following steps: Step one: computer reads data; use the computer to read the unlabeled SA-1B data set and labeled remote sensing image data, and respectively pre-process the image data; Step two: construct a state space model-based knowledge distillation network and a lightweight remote sensing target instance segmentation network; The state space model-based knowledge distillation network is composed of two parts of teacher network and student network, and the lightweight remote sensing target instance segmentation network is composed of three parts of backbone, neck and head; wherein the teacher network is the image encoder of the SAM base model, and no parameter adjustment is made during training; the backbone network is the student network obtained after training the state space model-based 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 three: train the neural network; after constructing the state space model-based knowledge distillation network using the Pytorch deep learning framework, train this network using unlabeled data, save the training parameter results, import the student network of the state space model-based knowledge distillation network, and then connect the neck and head to form the lightweight remote sensing target instance segmentation network; the specific training method is as follows: in the first training, the backbone is constructed by the state space model, then it is used as a student model to train the lightweight method using knowledge distillation, so that it learns the complex knowledge of the teacher model, then the training results of the backbone are reserved 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 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; Step four: remote sensing image instance segmentation, use the lightweight remote sensing target instance segmentation network model trained in step three to realize high-quality remote sensing image instance segmentation on the test set or actual image.
2. The method of claim 1, wherein the method is characterized in that: In step one, the data sources used are the SA-1B data set of SAM official and the famous public remote sensing data sets SSDD, WHU and NWPU data set; in the SA-1B data set, 1.2w pictures are selected and expanded to 5w 1024x1024 images after data augmentation; SSDD, WHU and NWPU data sets are small and medium-sized data sets, and the three are respectively for optical SAR ship, building static target and multi-type target mixed scene; select the complete data set and convert it into COCO format to adapt to the MmCV framework, also perform data augmentation, and expand to 1024x1024 size without changing the aspect ratio.
3. The method of claim 1, wherein the method is characterized in that: In step two, three different backbone networks are designed. The original Mamba, Vanilla Mamba, is composed of 12 layers of normal visual Mamba blocks, Vim Block. The model is the smallest, but the learning ability is weak. The Trans Mamba is composed of 12 layers of visual Mamba blocks combined with Transfromer, TRVim Block. Specifically, it contains five components, residual layer, normalization layer, Mamba layer, Gelu activation function layer and multi-layer perception, MLP. The learning ability is strong, but the model is the largest and the inference is slow. The Scanning Mamba is composed of 24 layers of scanning Mamba blocks, SCVim Block. The model size is in the middle, but the learning ability is strong and the inference is fast. Users can choose according to task requirements.
4. The method of claim 3, wherein the method is characterized in that: The three backbone networks all cut each input picture into J = n 2 p x p x c slices, flatten them, and multiply them by the training matrix W and the sine position encoding, respectively, to obtain the matrix X position is the sine position encoding. If X1is defined as the output result of the first layer, 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; and 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 ); 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 link 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 using 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 of claim 4, wherein the method is characterized in that: Neck is the connection structure between the backbone network and the head, when the recommended scan of the manba layer is selected, the odd layer X 2k-1 Output intermediate features, X 2k-1 represents the output result of the 2k-1 layer, Y 2k-1 = Downconv(X 2k-1 ), Y 2k-1 is the result of down-sampling, k is an integer between 1 and 12, and then output through a full convolution layer, which is represented as Z 2k+1 =Y 2k-1 +Z 2k-1 +conv(Y 2k-- +Z 2k-1 ), Z 2k+1 represents the intermediate feature obtained after the kth recursion, and finally the result of the 23th layer is simply feature pyramid to obtain F feature =[pooling(Z), Z, transpose(Z), transpose 2 (Z)], wherein Z represents the output of the last odd layer, pooling represents pooling, and transpose represents deconvolution.
6. The method of claim 5, wherein the method is characterized in that: When the other two layers of the mamba are selected, use the odd layer X 2p-1 The output intermediate feature, X 2p-1 represents the output result of the 2p-1 layer, and M 2p-1 = Downconv(X 2k-1 ), M 2p-1 is the result of down-sampling, and p is an integer between 1 and 6, and then output through a full convolution layer, specifically represented as N 2p+1 =M 2p-1 +N 2p-1 +conv(M 2p-1 +N 2p-1 ), N 2P+1 represents the intermediate feature obtained after the Pth recursion, and finally the result of the 11th layer is simply a feature pyramid, F feature =[pooling(Z), Z, transpose(Z), transpose 2 (Z)], wherein Z represents the output of the last odd layer, pooling represents pooling, and transpose represents deconvolution.
7. The method of claim 1, wherein the method is characterized in that: In step three, under the Pytorch deep learning framework, during the first training, the backbone is constructed by using the state space model as the student model to train the lightweight method of knowledge distillation. For the augmented 5w pictures, the teacher and student models are fed in, and the output results are both 4-dimensional matrices. The first dimension is flattened and the one-norm loss is directly calculated for parameter backpropagation, and the final parameters are retained.
8. The method 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-training 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, a part of the training set is selected for each epoch, and the order is randomly arranged. Every 3 rounds, the optimal network parameters are selected according to the test data. Finally, the optimal network parameters of the whole process are retained as the final training results.
9. The method of claim 8, wherein the method is characterized in that: During training, the AmpOptimWrapper optimization encapsulator type of mmdection is used to operate data of the bfloat16 floating-point type; the AdamW optimizer is used, and the learning rate base_lr and the weight decay rate 0.05 are set; the linear learning rate scheduler LinearLR and the cosine annealing learning rate scheduler CosineAnnealingLR are configured, the former starts from the iteration number 0 and lasts to the iteration number 50, the initial learning rate factor is 0.001, and the scheduling is based on the iteration number rather than the epoch number, the latter starts from the first epoch and lasts to the end of the maximum round max_epochs; the minimum learning rate eta_min is base_lr*0.001, and the cosine annealing period T_max is also set to max_epochs, and the scheduling is based on the epoch number, the learning rate is dynamically adjusted according to the iteration or epoch number, and the training process of the model is optimized; in addition, the images and the corresponding annotations are randomly horizontally flipped with a probability of 50%, and the RandomResize function is used to randomly adjust the size of the images, the scale changes based on the crop_size, and the aspect ratio changes within (0.1, 2.0); the filter FilterAnnotations is responsible for filtering out the annotation information whose bounding box or mask size is less than the given threshold 1e-5.
10. The method of claim 1, wherein the method is characterized in that: In step four, the lightweight remote sensing instance segmentation network model based on state space knowledge distillation trained in step three is used for instance segmentation in a remote sensing scene; when the user imports the parameters trained in step three into the lightweight remote sensing target instance segmentation network, directly runs the official test file image_demo.py, and inputs the image path that needs to be instance segmented into the model, the model uses the lightweight state space model backbone to obtain features, then uses the Neck feature fusion, uses the Mask-RCNN-Head to predict the final result, and the model outputs the visual result of the instance segmentation of the image in the same position, and a yaml type file storing the mask and bounding box data; on the SSDD, WHU and NWPU datasets, the target size scale of each test set image is tested, including ship and aircraft targets; when testing the mAP related indicators according to the standard dataset matched during training, directly run the official test file test.py, and under this framework, the model will automatically process each image in the test set and calculate the mAP indicator.
Citation Information
Patent Citations
Lung disease lesion unsupervised segmentation method based on knowledge distillation
CN113902761A
Image instance segmentation method and system, device and nonvolatile readable storage medium
WO2024113782A1