Two-stage knowledge distillation method based on dynamic sampling prompt

Abstract

The invention discloses a two-stage knowledge distillation method based on dynamic sampling prompt. The two-stage knowledge distillation method comprises the following steps of image encoder distillation and decoder distillation. Based on the design of two-stage knowledge distillation, compression and light weight of a large model are realized, synchronous distillation is carried out on an SAM image encoder and a decoder, efficient alignment of encoder extraction features and a mask decoder is realized, and the calculation cost is effectively reduced while the excellent generalization ability of the model is reserved. According to the method, a dynamic sampling prompt strategy is designed in decoder knowledge distillation and is used for efficiently selecting prompt points, so that the model learning process is optimized. The strategy particularly focuses on incorrect regions in the segmentation result to guide the student model to be closer to a true value. By dynamically updating the prompt point set, knowledge learning of the student model to the teacher model is gradually enhanced, and performance loss possibly caused in the distillation process is effectively reduced.

Description

Technical Field

[0001] The present invention relates to the technical field of knowledge distillation, and in particular to a two-stage knowledge distillation method based on dynamic sampling prompts. Background Art

[0002] As a new paradigm of generative artificial intelligence, large-scale visual models have opened up a new era for visual basic models. Among them, the SAM (Segment Anything Model) model, relying on the pre-training foundation on ultra-large-scale datasets, combined with cross-modal cue learning, is closely connected with downstream tasks, showing strong robustness and generalization ability. However, such large models usually consume a lot of computing resources and are difficult to apply to resource-constrained edge environments. Knowledge distillation is a model compression technology that aims to transfer the knowledge and representation capabilities of complex pre-trained large models (i.e., teacher models) to smaller and more efficient models (i.e., student models). Early knowledge distillation methods were mainly applied to classification tasks, but with the development of technology, this method has gradually been extended to dense prediction tasks such as object detection and semantic segmentation. In the study of knowledge distillation for SAM large models, existing methods usually focus on extracting the knowledge of SAM image encoders into lightweight student models, thereby significantly improving the inference speed. However, this type of research method has certain limitations. Since it only focuses on the distillation of image encoders and ignores the coordinated optimization between encoders and mask decoders, the features extracted by the encoder cannot be well aligned with the mask decoder, resulting in performance degradation. Summary of the invention

[0003] In order to solve the above problems in the prior art, the present invention proposes a two-stage knowledge distillation method based on dynamic sampling prompts, which can lightweight the SAM large model while retaining its excellent generalization performance.

[0004] The basic idea of ​​the present invention is: first, the present invention uses knowledge distillation technology to transfer the complex feature extraction capabilities of the heavyweight image encoder based on ViT (Vision Transformer) in the SAM large model to a lightweight image encoder. Subsequently, the lightweight image encoder generated by distillation is directly combined with the mask decoder of the SAM model to form a student model of decoder knowledge distillation. On this basis, the mask decoder is further optimized by distillation using the dynamic sampling prompt method to achieve efficient alignment of the features extracted by the encoder and the mask decoder.

[0005] In order to achieve the above object, the technical solution of the present invention is as follows: a two-stage knowledge distillation method based on dynamic sampling prompts, comprising the following steps:

[0006] A. Image Encoder Distillation

[0007] First, a lightweight backbone network is introduced to replace the image encoder in SAM, and the ViT-H model in SAM is replaced with the TinyViT model to construct a student model for encoder knowledge distillation. During the distillation process, the original ViT-H model is used as the teacher model to guide the TinyViT model to learn the knowledge representation of the teacher model. At the same time, the last downsampling operation of the TinyViT model is removed, and the projection layer is used to align the channel dimension. The specific encoder distillation loss is as follows:

[0008] L emb =L e (T enc (I),S enc (I))

[0009] Where I represents the input image, T enc and S enc denote the image encoders of the teacher model and the student model respectively. L e is the loss function, here we use mean square error loss. L emb is the encoder distillation loss function. During the forward propagation of the teacher model, the results of the data augmentation operation of the teacher model and the corresponding teacher model inference results are stored in advance, and these data are loaded during the training process to replicate the forward propagation process of the teacher model.

[0010] B. Decoder Distillation

[0011] The lightweight image encoder obtained by distillation in step A is directly combined with the mask decoder of the SAM model to construct a student model for decoder knowledge distillation. In this process, the network structure of the SAM decoder remains completely unchanged. In the model training stage, the traditional sample label is not used as a supervisory signal. Instead, the mask output by the SAM model is used as the true value to guide the student model of the decoder knowledge distillation to generate a mask output. The specific decoder distillation loss is as follows:

[0012] L mask =L d (T dec (T enc (I),q),S dec (S ec (I),q))

[0013] Among them, T dec and S dec Denote the mask decoders of the teacher model and the student model, T enc (I) and S enc (I) represents the image embedding of the teacher model and the student model. q represents the query of the mask decoder, including the hint embedding and the output token. Ld It is a combination of Dice loss and binary cross entropy loss. mask Distillation loss function for the decoder.

[0014] In the process of selecting the point cues that drive the decoder to output the predicted mask, a dynamic sampling cue strategy is used to select cue points. The specific steps are as follows: First, the initial cue points are sampled from the mask predicted by the teacher model, and these initial cue points are input into the student model to generate the output mask. Then, the difference areas between the teacher mask and the student mask are calculated, and new cue points are sampled in these areas. For the over-segmented areas of the student mask, negative cue points are sampled; for the under-segmented areas, positive cue points are sampled to guide the student model to output a mask closer to the true value. Subsequently, the newly sampled cue points are added to the existing cue point set, and the cue point set of each training sample is dynamically updated. For each original training sample in the student model, the updated cue point set is used to re-infer, and the decoder distillation loss between the student model output mask and the teacher model mask is calculated. The number of cue point updates and re-inferences for the same training sample is controlled by a preset threshold of iterations. Through this interactive iterative optimization process, the cue point set is continuously updated and the inference results of the student model are improved.

[0015] Furthermore, the image encoder distillation is trained using a 0.9%-1.2% subset of SAM's original training dataset SA-1B, which is a large-scale segmentation dataset released by Facebook's parent company MetaAI laboratory; the decoder distillation uses the same image data as the image encoder distillation, combines the image encoder obtained by image encoder distillation with SAM's mask decoder, initializes the decoder with SAM's weights, and uses a dynamic sampling prompt strategy to select prompt points.

[0016] Compared with the prior art, the present invention has the following beneficial effects:

[0017] 1. Based on the design of two-stage knowledge distillation, the present invention realizes the compression and lightweight of large models, performs synchronous distillation on the SAM image encoder and decoder, realizes the efficient alignment of encoder extracted features and mask decoder, and effectively reduces the computational cost while retaining the excellent generalization ability of the model.

[0018] 2. The present invention designs a dynamic sampling hint strategy in the decoder knowledge distillation to efficiently select hint points, thereby optimizing the model learning process. This strategy pays special attention to incorrect areas in the segmentation results to guide the student model closer to the true value. By dynamically updating the hint point set, the student model gradually strengthens the knowledge learning of the teacher model, effectively reducing the performance loss that may occur during the distillation process. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 It is a schematic diagram of the process of the present invention.

[0020] Figure 2 It is a schematic diagram of the dynamic sampling prompt strategy. DETAILED DESCRIPTION

[0021] The specific implementation of the present invention is described in detail below with reference to the accompanying drawings. The present invention proposes a two-stage knowledge distillation method based on dynamic sampling prompts, such as Figure 1 As shown in the figure, the overall process of this method includes two stages: first, the first stage image encoder distillation, in which a 1% subset of SAM's original training dataset SA-1B is used for training. SA-1B dataset is a large-scale segmentation dataset released by MetaAI Lab, Facebook's parent company; followed by the second stage decoder distillation, which still uses the same image data as the first stage, combines the image encoder obtained by the first stage distillation with SAM's mask decoder, and initializes the decoder with SAM's weights, and uses a dynamic sampling prompt strategy to select prompt points. The specific steps include the following:

[0022] A. First-stage image encoder distillation

[0023] First, a lightweight backbone network is introduced to replace the image encoder in SAM. The present invention replaces the ViT-H model in SAM with the TinyViT model to construct a student model for encoder knowledge distillation. During the distillation process, the original ViT-H model is used as the teacher model to guide the TinyViT model to learn the knowledge representation of the teacher model. At the same time, the last downsampling operation of the TinyViT model is removed, and the projection layer is used to align the channel dimension. The specific encoder distillation loss is as follows:

[0024] L emb =L e (T enc (I),S enc (I))

[0025] Where I represents the input image, T enc and S enc denote the image encoders of the teacher model and the student model respectively. L e is the loss function, here we use mean square error loss. L embis the encoder distillation loss function. During the forward propagation of the teacher model, the results of the data augmentation operation of the teacher model and the corresponding teacher model inference results are stored in advance, and these data are loaded during the training process to replicate the forward propagation process of the teacher model. In this embodiment of the present invention, the student model of the encoder knowledge distillation specifically uses the TinyViT-5M model with a small number of parameters, and uses the Adam optimizer on the 1% SA-1B dataset with L emb As the distillation loss, the training is performed for 8 epochs, and the batch size and initial learning rate are set to 8 and 5e respectively. -2 .

[0026] B. Second stage decoder distillation

[0027] The lightweight image encoder obtained by distillation in step A is directly combined with the mask decoder of the SAM model to construct a student model for decoder knowledge distillation. In this process, the network structure of the SAM decoder remains completely unchanged. In the model training stage, the traditional sample label is not used as the supervision signal, but the mask output by the SAM model is used as the true value to guide the student model of the decoder knowledge distillation to generate the mask output. The specific decoder distillation loss is as follows:

[0028] L mask =L d (T dec (T enc (I),q),S dec (S enc (I),q))

[0029] Among them, T dec and S dec Denote the mask decoders of the teacher model and the student model, T enc (I) and S enc (I) represents the image embedding of the teacher model and the student model. q represents the query of the mask decoder, including the hint embedding and the output token. L d It is a combination of Dice loss and binary cross entropy loss. mask Distillation loss function for the decoder.

[0030] In the process of selecting the point hints that drive the decoder to output the predicted mask, the present invention designs a dynamic sampling hint strategy for efficiently selecting hint points. The specific steps are as follows: First, sample the initial hint points from the mask predicted by the teacher model, and input these initial hint points into the student model to generate the output mask. Then, calculate the difference area between the teacher mask and the student mask, and sample new hint points in these areas. Figure 2As shown, for the over-segmented area of ​​the student mask, negative cue points are sampled; for the under-segmented area, positive cue points are sampled to guide the student model to output a mask closer to the true value. Subsequently, the newly sampled cue points are added to the existing cue point set, and the cue point set of each training sample is dynamically updated. For each original training sample in the student model, the updated cue point set is used to re-infer, and the decoder distillation loss between the student model output mask and the teacher model mask is calculated. The number of cue point updates and re-inferences for the same training sample is controlled by a preset iteration number threshold. In this embodiment of the present invention, the number of iterations is set to 5. Through this interactive iterative optimization process, the cue point set is continuously updated and the inference results of the student model are improved, thereby achieving efficient distillation of the decoder.

[0031] In the embodiment of the present invention, the decoder knowledge distillation is also trained using a 1% subset of the SAM original training data set SA-1B, and L mask The distillation loss is used for training for 5 epochs, and the batch size and initial learning rate are set to 5 and 1e respectively. -4 .

[0032] The above describes the basic principles, main features and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the above embodiments, and the above embodiments and descriptions are only for explaining the principles of the present invention. Without departing from the spirit and scope of the present invention, the present invention may have various changes and improvements, and these changes and improvements fall within the scope of the present invention to be protected.

Claims

1. A two-stage knowledge distillation method based on dynamic sampling prompts, characterized by: The following steps are included: A. Image Encoder Distillation First, a lightweight backbone network is introduced to replace the image encoder in SAM, and the ViT-H model in SAM is replaced by the TinyViT model to construct a student model for encoder knowledge distillation. During the distillation process, the original ViT-H model is used as the teacher model to guide the TinyViT model to learn the knowledge representation of the teacher model. At the same time, the last downsampling operation of the TinyViT model is removed, and the projection layer is used to align the channel dimension. The specific encoder distillation loss is as follows: L emb =L e (T enc (I),S enc (I)) Where I represents the input image, T enc and S enc Denote the image encoders of the teacher model and the student model respectively; L e is the loss function, here we use the mean square error loss; L emb Distill the loss function for the encoder; during the forward propagation of the teacher model, store the results of the data augmentation operation of the teacher model and the corresponding teacher model inference results in advance, and load these data during the training process to replicate the forward propagation process of the teacher model; B. Decoder Distillation The lightweight image encoder obtained by distillation in step A is directly combined with the mask decoder of the SAM model to construct a student model for decoder knowledge distillation. In this process, the network structure of the SAM decoder remains completely unchanged. In the model training stage, the traditional sample label is not used as the supervision signal, but the mask output by the SAM model is used as the true value to guide the student model of the decoder knowledge distillation to generate the mask output. The specific decoder distillation loss is as follows: L mask =L d (T dec (T enc (I),q),S dec (S enc (I),q)) Among them, T dec and S dec Denote the mask decoders of the teacher model and the student model, T enc (I) and S enc (I) represents the image embedding of the teacher model and the student model; q represents the query of the mask decoder, including the hint information embedding and output token; L d It is a combination of Dice loss and binary cross entropy loss; L mask Distill loss function for the decoder; In the process of selecting the point cues that drive the decoder to output the predicted mask, a dynamic sampling cue strategy is used to select cue points. The specific steps are as follows: first, initial cue points are sampled from the mask predicted by the teacher model, and these initial cue points are input into the student model to generate the output mask. Then, the difference areas between the teacher mask and the student mask are calculated, and new cue points are sampled in these areas. For the over-segmented areas of the student mask, negative cue points are sampled. For the under-segmented areas, positive cue points are sampled to guide the student model to output a mask closer to the true value. Subsequently, the newly sampled cue points are added to the existing cue point set, and the cue point set of each training sample is dynamically updated. For each original training sample in the student model, the updated cue point set is used to re-infer, and the decoder distillation loss between the student model output mask and the teacher model mask is calculated. The number of cue point updates and re-inferences for the same training sample is controlled by a preset iteration threshold. Through this interactive iterative optimization process, the cue point set is continuously updated and the inference results of the student model are improved.

2. According to claim 1, a two-stage knowledge distillation method based on dynamic sampling prompts is characterized by: The image encoder distillation is trained using a 0.9%-1.2% subset of SAM's original training dataset SA-1B, which is a large-scale segmentation dataset released by Facebook's parent company MetaAI laboratory. The decoder distillation uses the same image data as the image encoder distillation, combines the image encoder obtained by image encoder distillation with SAM's mask decoder, and initializes the decoder with SAM's weights, and uses a dynamic sampling prompt strategy to select prompt points.

Images