Method and device for knowledge distillation of mask autoencoder, equipment and storage medium
By employing a knowledge distillation method based on a teacher-student framework, a small-scale student model is trained using a large-scale teacher model. This addresses the issue of large parameter quantities in masked autoencoders deployed at the edge of the power grid, thereby improving model accuracy and real-time inference speed.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- GLOBAL ENERGY INTERCONNECTION RES INST CO LTD
- Filing Date
- 2023-01-31
- Publication Date
- 2026-04-14
AI Technical Summary
Large-scale pre-trained masked autoencoder models have a large number of parameters, making them difficult to deploy at the edge of the power grid for real-time power image classification. Existing compression methods cannot effectively reduce model parameters and affect model accuracy.
A teacher-student framework is adopted to establish teacher and student models. The teacher model is a large-scale visual transformation model. Through pre-training and knowledge distillation techniques, the knowledge of the teacher model is distilled into the student model, and fine-tuning training is carried out in combination with downstream tasks to reduce model parameters while maintaining accuracy.
While reducing model parameters, the model accuracy is maintained, and the data generalization ability of the student model is improved, making it suitable for real-time inference on the power edge with limited computing resources.
Smart Images

Figure CN116227582B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the field of mask autoencoder compression technology, and in particular to a knowledge distillation method, apparatus, device, and storage medium for a mask autoencoder. Background Technology
[0002] Currently, large-scale pre-trained models perform very well in downstream tasks. However, in specific power grid applications, such as power transmission and transformation image defect detection, high real-time requirements necessitate the classification of power images at the edge of the power grid. Since edge computing power and storage space are limited, large models are not suitable for direct deployment. Therefore, compression of these models is necessary. A widely used model compression method is knowledge distillation based on a teacher-student framework. This method first pre-trains a large-parameter teacher model, then uses the intermediate or final output of the teacher model to supervise the training of the student model, distilling teacher knowledge into the student model to improve its learning ability and generalization capabilities.
[0003] In recent years, ViTs (including multiple Vision Transformer (ViT) models) that use Transformer models to divide images into a sequence of image patches have received a lot of attention. Pre-trained masked autoencoders (MAEs, model structures of ViTs) have performed well in downstream tasks. However, these models have a large number of parameters and there are no corresponding compression methods to simplify their model structure and reduce their number of parameters, making them difficult to deploy at the edge of the power grid and unable to quickly use field-collected data for fine-tuning for downstream tasks. Summary of the Invention
[0004] In view of this, embodiments of the present invention provide a knowledge distillation method, apparatus, device and storage medium for a mask autoencoder to solve the technical problem of large model parameters in a mask autoencoder.
[0005] The technical solution proposed in this invention is as follows:
[0006] The first aspect of this invention provides a knowledge distillation method for a mask autoencoder, comprising:
[0007] Teacher and student models of a mask autoencoder are established separately, wherein both the teacher model and the student model are visual transformation models, and the teacher model is larger in scale than the student model; the teacher model is pre-trained; the student model is pre-trained by knowledge distillation based on the pre-trained teacher model; and the pre-trained student model is fine-tuned based on downstream tasks.
[0008] Optionally, the pre-training of the teacher model includes: masking a first input image according to a first preset masking rate; inputting the masked first input image into the encoder of the teacher model to obtain a first intermediate feature; inputting the first intermediate feature into the decoder of the teacher model to obtain a first reconstructed image; obtaining a first loss function based on the first input image and the first reconstructed image; and pre-training the teacher model based on the first loss function.
[0009] Optionally, the step of performing knowledge distillation pre-training on the student model based on the pre-trained teacher model includes: obtaining reconstruction loss and distillation loss based on the student model and the pre-trained teacher model; determining a second loss function based on the reconstruction loss and distillation loss; and performing knowledge distillation pre-training on the student model based on the second loss function.
[0010] Optionally, obtaining the reconstruction loss and distillation loss based on the student model and the pre-trained teacher model includes: masking the second input image according to a second preset masking rate; inputting the masked second input image into the encoder of the student model to obtain a second intermediate feature, and inputting the second intermediate feature into the decoder of the student model to obtain a second reconstructed image; inputting the masked second input image into the encoder of the pre-trained teacher model to obtain a third intermediate feature, and inputting the third intermediate feature into the decoder of the pre-trained teacher model to obtain a third reconstructed image; obtaining the reconstruction loss based on the second reconstructed image; and obtaining the distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image.
[0011] Optionally, obtaining the reconstruction loss based on the second reconstructed image includes: obtaining the distance between the pixels of the second input image and the second reconstructed image in the masked region, and the distance between the pixels of the second input image and the second reconstructed image in the unmasked region; and obtaining the reconstruction loss based on the distance between the pixels in the masked region and the distance between the pixels in the unmasked region.
[0012] Optionally, obtaining the distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image includes: calculating the distance between the second intermediate feature and the third intermediate feature; calculating the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked region based on a preset temperature coefficient; and obtaining the distillation loss based on the distance between the second intermediate feature and the third intermediate feature, and the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked region.
[0013] Optionally, the step of fine-tuning the pre-trained student model based on the downstream task includes: obtaining the fourth intermediate feature of the training set of the downstream task according to the encoder of the pre-trained student model; and inputting the fourth intermediate feature into the mapping layer of the downstream task for fine-tuning training.
[0014] A second aspect of the present invention provides a knowledge distillation apparatus for a mask autoencoder, comprising:
[0015] The model building module is used to build a teacher model and a student model for the mask autoencoder, respectively. Both the teacher model and the student model are visual transformation models, and the teacher model is larger than the student model. The teacher model training module is used to pre-train the teacher model. The knowledge distillation pre-training module is used to pre-train the student model based on the pre-trained teacher model. The fine-tuning training module is used to fine-tune the pre-trained student model based on downstream tasks.
[0016] Optionally, the teacher model training module includes: a first masking module for masking a first input image according to a first preset masking rate; a first input module for inputting the masked first input image into the encoder of the teacher model to obtain a first intermediate feature; a first reconstruction module for inputting the first intermediate feature into the decoder of the teacher model to obtain a first reconstructed image; a first loss module for obtaining a first loss function based on the first input image and the first reconstructed image; and a first pre-training module for pre-training the teacher model based on the first loss function.
[0017] Optionally, the knowledge distillation pre-training module includes: a first acquisition module for acquiring reconstruction loss and distillation loss based on the student model and the pre-trained teacher model; a second loss module for determining a second loss function based on the reconstruction loss and distillation loss; and a second pre-training module for performing knowledge distillation pre-training on the student model based on the second loss function.
[0018] Optionally, the first acquisition module includes: a second masking module, used to mask the second input image according to a second preset masking rate; a second reconstruction module, used to input the masked second input image into the encoder of the student model to obtain second intermediate features, and input the second intermediate features into the decoder of the student model to obtain a second reconstructed image; a third reconstruction module, used to input the masked second input image into the encoder of a pre-trained teacher model to obtain third intermediate features, and input the third intermediate features into the decoder of a pre-trained teacher model to obtain a third reconstructed image; a second loss module, used to obtain a reconstruction loss based on the second reconstructed image; and a third loss module, used to obtain a distillation loss based on the second intermediate features, the third intermediate features, the second reconstructed image, and the third reconstructed image.
[0019] Optionally, the second loss module includes: a second acquisition module, used to acquire the distance between the second input image and the second reconstructed image in the masked region, and the distance between the second input image and the second reconstructed image in the unmasked region; and a reconstruction loss module, used to acquire a reconstruction loss based on the distance between the pixels in the masked region and the distance between the pixels in the unmasked region.
[0020] Optionally, the third loss module includes: a calculation module for calculating the distance between the second intermediate feature and the third intermediate feature; a distance module for calculating the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked area based on a preset temperature coefficient; and a distillation loss module for obtaining distillation loss based on the distance between the second intermediate feature and the third intermediate feature, and the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked area.
[0021] Optionally, the fine-tuning training module includes: a third acquisition module, used to acquire a fourth intermediate feature of the training set of the downstream task based on the encoder of the pre-trained student model; and a fine-tuning module, used to input the fourth intermediate feature into the mapping layer of the downstream task for fine-tuning training. A third aspect of the present invention provides an electronic device, including: a memory and a processor, the memory and the processor being communicatively connected to each other, the memory storing computer instructions, and the processor executing the computer instructions to perform the knowledge distillation method of the mask autoencoder as described in any one of the first aspects of the present invention.
[0022] A fourth aspect of the present invention provides a computer-readable storage medium storing computer instructions for causing the computer to perform a knowledge distillation method of a mask autoencoder as described in any of the first aspects of the present invention.
[0023] As can be seen from the above technical solutions, the embodiments of the present invention have the following advantages:
[0024] This invention provides a knowledge distillation method, apparatus, device, and storage medium for a masked autoencoder. It establishes a teacher model and a student model for the masked autoencoder, both of which are visual transformation models, with the teacher model being larger than the student model. The teacher model is pre-trained. Based on the pre-trained teacher model, the student model undergoes knowledge distillation pre-training, enabling the student model to learn data generalization ability from the pre-trained teacher model and obtain image features with better representation capabilities. The pre-trained student model is then fine-tuned based on downstream tasks. The student model can be deployed on power edge devices with limited computing resources, reducing model parameters while maintaining model accuracy and accelerating real-time inference speed. Attached Figure Description
[0025] To more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.
[0026] Figure 1 This is a flowchart of the knowledge distillation method of the mask autoencoder in an embodiment of the present invention;
[0027] Figure 2 This is an architecture diagram of the knowledge distillation method of the mask autoencoder in an embodiment of the present invention;
[0028] Figure 3 This is a schematic diagram of the knowledge distillation device of the mask autoencoder in an embodiment of the present invention;
[0029] Figure 4 This is a schematic diagram of the structure of the electronic device in an embodiment of the present invention;
[0030] Figure 5 This is a schematic diagram of the structure of a computer-readable storage medium in an embodiment of the present invention. Detailed Implementation
[0031] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.
[0032] This invention provides a knowledge distillation method for a mask autoencoder, such as... Figure 1 and Figure 2 As shown, the method includes:
[0033] Step S100: Establish teacher and student models for the masked autoencoder, respectively. Both the teacher and student models are Vision Transformer (ViT) models, with the teacher model being larger than the student model. Specifically, both the teacher and student models have a ViT network structure, including an encoder and a decoder. The process is as follows: the masked image is input into the encoder to obtain intermediate features, and then these intermediate features are input into the decoder to reconstruct the image. The difference between the teacher and student models is that the teacher model is larger, meaning it has more parameters and a higher network complexity.
[0034] Step S200: Pre-train the teacher model. By pre-training the teacher model, which has relatively many parameters and a complex network structure, a pre-trained teacher model is obtained. This pre-trained teacher model provides supervised parameters based on knowledge distillation for the pre-training of the student model.
[0035] Step S300: Pre-train the student model using knowledge distillation based on the pre-trained teacher model. During the pre-training of the student model, the input image is synchronously input into the teacher model, and the information from the teacher model is used to supervise the student model. By supervising the student model with information from the teacher model, the student model can learn data generalization ability from the pre-trained teacher model, obtaining image features with better representation capabilities. This reduces model parameters during the pre-training stage while maintaining model accuracy, thus accelerating real-time inference speed.
[0036] Step S400: Fine-tune the pre-trained student model based on the downstream task. Fine-tuning is to retrain the pre-trained student model using the downstream task dataset to improve its performance on the downstream task dataset. For example, fine-tune the pre-trained student model using the power image dataset for the image classification task.
[0037] This invention discloses a knowledge distillation method for a masked autoencoder. This method involves establishing a teacher model and a student model for the masked autoencoder, both of which are visual transformation models, with the teacher model being larger than the student model. The teacher model is pre-trained. Based on the pre-trained teacher model, the student model undergoes knowledge distillation pre-training. Finally, the pre-trained student model is fine-tuned based on downstream tasks. This allows the student model to learn data generalization ability from the pre-trained teacher model, resulting in image features with better representation capabilities. The student model can be deployed on power edge devices with limited computing resources, reducing model parameters while maintaining model accuracy and accelerating real-time inference.
[0038] In one embodiment, step S200, which pre-trains the teacher model, includes: masking a first input image according to a first preset masking rate; inputting the masked first input image into the encoder of the teacher model to obtain a first intermediate feature; inputting the first intermediate feature into the decoder of the teacher model to obtain a first reconstructed image; obtaining a first loss function based on the first input image and the first reconstructed image; and pre-training the teacher model based on the first loss function.
[0039] Specifically, the first input image x is an electricity image from a large-scale dataset. It is first divided into N equal-sized image patches. A predetermined first mask ratio (MR) is used to randomly mask the image. Selected patches are invisible to the teacher model. The remaining visible patches are concatenated into a sequence, which is then input into the encoder to obtain the first intermediate features. The decoder uses these first intermediate features and a mask vector randomly obtained according to the first mask ratio to decode and obtain the first reconstructed image x. ′ Calculate the first input image x and the first reconstructed image x in pixel space. ′ The mean squared error (MSE) between the two models is used as the first loss function to calculate a loss value. Then, the gradient is calculated in reverse to update the parameters of the teacher model, thereby completing the pre-training of the teacher model.
[0040] In one embodiment, step S300, which involves pre-training the student model using knowledge distillation based on a pre-trained teacher model, includes:
[0041] Step S310: Obtain reconstruction loss and distillation loss based on the student model and the pre-trained teacher model;
[0042] Step S320: Determine the second loss function based on the reconstruction loss and distillation loss;
[0043] Step S330: Perform knowledge distillation pre-training on the student model based on the second loss function.
[0044] Specifically, during the pre-training of the student model, the input image is synchronously input into the teacher model. The information from the teacher model is used to supervise the student model. The distillation loss is obtained based on the information from the teacher model, while the reconstruction loss is calculated by the student model based on its own input image and the reconstructed image. The second loss function is jointly determined by the reconstruction loss and the distillation loss. The influence of these two losses on the student model is determined by the hyperparameter β. The hyperparameter β can be selected based on the best-performing value through experiments. The second loss function L during the pre-training phase of the student model... s As shown in the following formula:
[0045] L s = recons +L kd
[0046] In the formula, L recons L represents the reconstruction loss. kd To represent distillation loss, the second loss function L is used. s A loss value is calculated, and then the gradient is calculated in reverse to update the parameters of the student model, thereby completing the pre-training of the student model. The student model can learn the data generalization ability from the pre-trained teacher model and obtain image features with better representation ability. The model accuracy is not reduced while reducing the model parameters in the pre-training stage, thus accelerating the real-time inference speed.
[0047] In one embodiment, step S330 above, which obtains reconstruction loss and distillation loss based on the student model and the pre-trained teacher model, includes:
[0048] Step S331: Mask the second input image according to the second preset mask rate;
[0049] Step S332: Input the masked second input image into the encoder of the student model to obtain the second intermediate features, and input the second intermediate features into the decoder of the student model to obtain the second reconstructed image;
[0050] Step S333: Input the masked second input image into the encoder of the pre-trained teacher model to obtain the third intermediate feature, and input the third intermediate feature into the decoder of the pre-trained teacher model to obtain the third reconstructed image;
[0051] Step S334: Obtain the reconstruction loss based on the second reconstructed image;
[0052] Step S335: Obtain the distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image.
[0053] Specifically, similar to the training process of the teacher model, the second input image X is an electric power image from a large-scale dataset. It is first divided into N equal-sized image patches. The image is then randomly masked using a predetermined second masking rate. In this embodiment, the first and second masking rates are the same, and the first and second input images are selected from the same dataset. The selected patches are invisible to both the teacher and student models. The remaining visible patches are concatenated into a sequence and synchronously input into the encoder of the student model and the encoder of the pre-trained teacher model. The encoder of the student model obtains the second intermediate feature f based on the masked second input image X. s , the second intermediate feature f s The second reconstructed image X is obtained by inputting the mask vector, which is randomly obtained according to the second mask rate, into the encoder of the student model. ′ The encoder of the pre-trained teacher model obtains the third intermediate feature f based on the masked second input image X. t The third intermediate feature f t The third reconstructed image t is obtained by inputting the mask vector, which is randomly obtained according to the second mask rate, into the encoder of the pre-trained teacher model. ′ .
[0054] Reconstruction loss and distillation loss are used together to determine the second loss function of the student model, based on the second reconstructed image X. ′ The reconstruction loss is determined by the variance of the pixels in the second input image X, and this reconstruction loss is based on the student model's own pre-training. This is based on the second intermediate feature f. s Third intermediate feature f t and the third reconstructed image t ′ The distillation loss is obtained, which provides intermediate features and soft label supervision for the pre-training of the student model.
[0055] In this embodiment of the invention, the student model masks the second input image during the pre-training phase and supervises the student model based on the pre-trained teacher model. This eliminates the need for labeling large-scale datasets of power grid scene images containing the second input image, reducing the manpower cost of labeling and allowing for rapid model updates using field-collected samples, thus reducing model iteration costs. Simultaneously, by incorporating knowledge distillation techniques, the number of student model parameters is reduced, improving the student model's data generalization ability.
[0056] In one embodiment, step S334, obtaining the reconstruction loss based on the second reconstructed image, includes: obtaining the distance between the pixels of the second input image and the second reconstructed image in the masked region, and the distance between the pixels of the second input image and the second reconstructed image in the unmasked region; and obtaining the reconstruction loss based on the distance between the pixels in the masked region and the distance between the pixels in the unmasked region.
[0057] Specifically, the reconstruction loss is shown in the following equation:
[0058]
[0059] In the formula, N m N u represents the number of pixels in the masked area and the number of pixels in the unmasked area, respectively. For masked pixels, since they are not input into the model, the closer the reconstructed pixels are to the real pixels, the better the model's ability to predict unknown pixels. Therefore, the reconstruction loss weight of unmasked pixels is reduced, that is, the right side of the plus sign in the above formula is multiplied by a hyperparameter c less than 1.
[0060] The embodiments of the present invention employ a weighted method that considers the combined influence of masked and unmasked regions on image reconstruction. Unlike existing technologies that only calculate the masked region, this method enhances the student's ability to reconstruct the unmasked region, and the decoder prevents over-focusing on the masked region while ignoring information from the unmasked region.
[0061] In one embodiment, step S335, obtaining distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image, includes: calculating the distance between the second intermediate feature and the third intermediate feature; calculating the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked area based on a preset temperature coefficient; and obtaining the distillation loss based on the distance between the second intermediate feature and the third intermediate feature, and the distance between the pixels of the second reconstructed image and the third reconstructed image in the masked area.
[0062] Specifically, the distillation loss is shown in the following formula:
[0063]
[0064] In the formula, The function is used to calculate the second intermediate feature f. s and the third intermediate feature f t The distance between them is calculated using the L1 paradigm, i.e., L1(,b)=||-b||1. The right half of the equation is used to calculate the second reconstructed image X of the student model. ′ The masked pixels and the third reconstructed image t of the teacher model ′ The distance between the masked pixels, where T is a preset temperature coefficient that controls the smoothness. The soft labels provided by the teacher model, with T values often greater than 1, mitigate the bias in the predictions of the pre-trained teacher model. Distillation loss addresses the lack of class labels for supervision during the pre-training of masked autoencoders. It utilizes intermediate features encoded by the teacher model and the reconstructed pixel predictions to provide supervision for the student model. In other words, distillation loss provides intermediate features and soft label supervision for the pre-training of the student model, enabling the student model to learn its ability to represent data and improve generalization.
[0065] In one embodiment, step S400, which involves fine-tuning the pre-trained student model based on the downstream task, includes: obtaining the fourth intermediate feature of the training set of the downstream task according to the encoder of the pre-trained student model; and inputting the fourth intermediate feature into the mapping layer of the downstream task for fine-tuning training.
[0066] Specifically, fine-tuning training involves retraining a pre-trained model using a downstream task training set to improve its performance on downstream tasks.
[0067] Unlike the pre-training phase, fine-tuning training only uses the encoder part of the student model. The training set for the downstream task is labeled, and the images from this set are input into the pre-trained student model's encoder to obtain the fourth intermediate feature. This fourth intermediate feature is then input into the mapping layer of the downstream task for fine-tuning. The mapping layer is a network added to the pre-trained student model based on the specific downstream task, used to perform that task. For example, if the downstream task is classification, only the fourth intermediate feature needs to be input into a fully connected layer and a softmax layer to obtain the classification result. The classification loss is then calculated, and fine-tuning is performed based on this loss. After fine-tuning, the student model can be used for inference. For example, for an image sample collected on-site, it can be input into the student model, and the output will be a classification result, classifying it as a certain type of defect or the absence of a defect.
[0068] In practical applications, the model can be quickly updated simply by fine-tuning the pre-trained student model according to the specific downstream task, thus reducing the cost of model iteration.
[0069] This invention also provides a knowledge distillation apparatus for a mask autoencoder, such as... Figure 3 As shown, the device includes:
[0070] The model building module 310 is used to build the teacher model and student model of the mask autoencoder respectively. Both the teacher model and the student model are visual transformation models, and the teacher model is larger in scale than the student model. For details, please refer to the corresponding part of the above method embodiment, which will not be repeated here.
[0071] The teacher model training module 320 is used to pre-train the teacher model; for details, please refer to the corresponding part of the above method embodiment, which will not be repeated here.
[0072] The knowledge distillation pre-training module 330 is used to perform knowledge distillation pre-training on the student model based on the pre-trained teacher model; for details, please refer to the corresponding part of the above method embodiment, which will not be repeated here.
[0073] The fine-tuning training module 340 is used to fine-tune the pre-trained student model based on downstream tasks. For details, please refer to the corresponding sections of the above method embodiments, which will not be repeated here.
[0074] This invention discloses a knowledge distillation device for a masked autoencoder. The device establishes a teacher model and a student model for the masked autoencoder, both of which are visual transformation models, with the teacher model being larger than the student model. The teacher model is pre-trained. Based on the pre-trained teacher model, the student model undergoes knowledge distillation pre-training. The pre-trained student model is then fine-tuned based on downstream tasks. This allows the student model to learn data generalization ability from the pre-trained teacher model, resulting in image features with better representation capabilities. The student model can be deployed on power edge devices with limited computing resources, reducing model parameters while maintaining model accuracy and accelerating real-time inference speed.
[0075] In one embodiment, the teacher model training module 320 includes:
[0076] The first masking module is used to mask the first input image according to a first preset masking rate;
[0077] The first input module is used to input the masked first input image into the encoder of the teacher model to obtain the first intermediate feature;
[0078] The first reconstruction module is used to input the first intermediate features into the decoder of the teacher model to obtain the first reconstructed image;
[0079] The first loss module is used to obtain a first loss function based on the first input image and the first reconstructed image;
[0080] The first pre-training module is used to pre-train the teacher model based on the first loss function.
[0081] In one embodiment, the knowledge distillation pre-training module 330 includes:
[0082] The first acquisition module is used to acquire reconstruction loss and distillation loss based on the student model and the pre-trained teacher model;
[0083] The second loss module is used to determine the second loss function based on the reconstruction loss and distillation loss;
[0084] The second pre-training module is used to pre-train the student model using knowledge distillation based on the second loss function.
[0085] In one embodiment, the first acquisition module includes:
[0086] The second masking module is used to mask the second input image according to the second preset masking rate;
[0087] The second reconstruction module is used to input the masked second input image into the encoder of the student model to obtain the second intermediate features, and input the second intermediate features into the decoder of the student model to obtain the second reconstructed image;
[0088] The third reconstruction module is used to input the masked second input image into the encoder of the pre-trained teacher model to obtain the third intermediate feature, and input the third intermediate feature into the decoder of the pre-trained teacher model to obtain the third reconstructed image.
[0089] The second loss module is used to obtain the reconstruction loss based on the second reconstructed image;
[0090] The third loss module is used to obtain the distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image.
[0091] In one embodiment, the second loss module includes:
[0092] The second acquisition module is used to acquire the distance between the pixels of the second input image and the second reconstructed image in the masked area, and the distance between the pixels of the second input image and the second reconstructed image in the unmasked area;
[0093] The reconstruction loss module is used to obtain the reconstruction loss based on the distance between pixels in the masked area and the distance between pixels in the unmasked area.
[0094] In one embodiment, the third loss module includes:
[0095] The calculation module is used to calculate the distance between the second and third intermediate features;
[0096] The distance module is used to calculate the distance between pixels in the masked area of the second and third reconstructed images based on a preset temperature coefficient.
[0097] The distillation loss module is used to obtain the distillation loss based on the distance between the second and third intermediate features, and the distance between the pixels of the second and third reconstructed images in the masked region.
[0098] In one embodiment, the fine-tuning training module 340 includes:
[0099] The third acquisition module is used to acquire the fourth intermediate feature of the training set of the downstream task based on the encoder of the pre-trained student model.
[0100] The fine-tuning module is used to input the fourth intermediate feature into the mapping layer of the downstream task for fine-tuning training.
[0101] This invention also provides an electronic device, such as... Figure 4As shown, the system includes a memory 420 and a processor 410, which are interconnected. The memory 420 stores computer instructions, and the processor 410 executes these computer instructions to perform the knowledge distillation method of the mask autoencoder as described in the above embodiments of the present invention. The processor 410 and the memory 420 can be connected via a bus or other means. The processor 410 can be a central processing unit (CPU). The processor 410 can also be other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, or combinations thereof. The memory 420, as a non-transitory computer storage medium, can be used to store non-transitory software programs, non-transitory executable programs, and modules, such as the corresponding program instructions / modules in the embodiments of the present invention. Processor 410 executes various functional applications and data processing by running non-transitory software programs, instructions, and modules stored in memory 420, thereby implementing the knowledge distillation method of the mask autoencoder in the above method embodiments. Memory 420 may include a program storage area and a data storage area. The program storage area may store applications required for operating the device and at least one function; the data storage area may store data created by processor 410, etc. Furthermore, memory 420 may include high-speed random access memory 420, and may also include non-transitory memory 420, such as at least one disk storage device, flash memory device, or other non-transitory solid-state storage device. In some embodiments, memory 420 may optionally include remotely located memories 420 relative to processor 410, which can be connected to processor 410 via a network. Examples of such networks include, but are not limited to, the Internet, intranets, local area networks, mobile communication networks, and combinations thereof. One or more modules are stored in memory 420 and, when executed by processor 410, perform the knowledge distillation method of the mask autoencoder as described in the above method embodiments. The specific details of the aforementioned electronic device can be understood by referring to the relevant descriptions and effects in the above method embodiments, and will not be repeated here.
[0102] This invention also provides a computer-readable storage medium, such as... Figure 5As shown, a computer program 510 is stored on the storage medium. When executed by a processor, this program implements the steps of the knowledge distillation method of the mask autoencoder in the above embodiments. The storage medium also stores audio and video stream data, feature frame data, interactive request signaling, encrypted data, and a preset data size. The storage medium can be a magnetic disk, optical disk, read-only memory (ROM), random access memory (RAM), flash memory, hard disk drive (HDD), or solid-state drive (SSD), etc.; the storage medium can also include combinations of the above types of memory. Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The computer program 13 can be stored in a computer-readable storage medium, and when executed, it can include the processes of the embodiments of the above methods. The storage medium can be a magnetic disk, optical disk, read-only memory (ROM), random access memory (RAM), flash memory, hard disk drive (HDD), or solid-state drive (SSD), etc.; the storage medium can also include a combination of the above types of memory.
[0103] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.
Claims
1. A knowledge distillation method for a mask autoencoder, characterized in that, include: Teacher and student models of a mask autoencoder are established separately, wherein both the teacher model and the student model are visual transformation models, and the teacher model is larger in size than the student model. The teacher model is pre-trained; The student model is pre-trained using knowledge distillation based on the pre-trained teacher model. The pre-trained student model is fine-tuned based on downstream tasks; The step of performing knowledge distillation pre-training on the student model based on the pre-trained teacher model includes: obtaining reconstruction loss and distillation loss based on the student model and the pre-trained teacher model; determining a second loss function based on the reconstruction loss and distillation loss; and performing knowledge distillation pre-training on the student model based on the second loss function. The step of obtaining reconstruction loss and distillation loss based on the student model and the pre-trained teacher model includes: masking the second input image according to a second preset masking rate; inputting the masked second input image into the encoder of the student model to obtain a second intermediate feature, and inputting the second intermediate feature into the decoder of the student model to obtain a second reconstructed image; inputting the masked second input image into the encoder of the pre-trained teacher model to obtain a third intermediate feature, and inputting the third intermediate feature into the decoder of the pre-trained teacher model to obtain a third reconstructed image; obtaining reconstruction loss based on the second reconstructed image; and obtaining distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image. The step of obtaining the reconstruction loss based on the second reconstructed image includes: obtaining the distance between the pixels of the second input image and the second reconstructed image in the masked region, and the distance between the pixels of the second input image and the second reconstructed image in the unmasked region; obtaining the reconstruction loss based on the distance between the pixels in the masked region and the distance between the pixels in the unmasked region, wherein the calculation formula for the reconstruction loss is: In the formula, Indicates the reconstruction loss. , These represent the number of pixels in the masked area and the number of pixels in the unmasked area, respectively. The value of the i-th pixel in the second input image. This represents the value of the i-th pixel in the second reconstructed image. The step of obtaining distillation loss based on the second intermediate feature, the third intermediate feature, the second reconstructed image, and the third reconstructed image includes: calculating the distance between the second intermediate feature and the third intermediate feature; calculating the distance between pixels in the masked region of the second reconstructed image and the third reconstructed image based on a preset temperature coefficient; and obtaining the distillation loss based on the distance between the second intermediate feature and the third intermediate feature, and the distance between pixels in the masked region of the second reconstructed image and the third reconstructed image. The formula for calculating the distillation loss is as follows: In the formula, The function is used to calculate the second intermediate feature. and the third intermediate feature The distance between them is determined by using Paradigm The third reconstructed image for the teacher model The value of the i-th pixel This is the preset temperature coefficient.
2. The knowledge distillation method for a mask autoencoder according to claim 1, characterized in that, The pre-training of the teacher model includes: The first input image is masked according to the first preset mask rate; The masked first input image is input into the encoder of the teacher model to obtain the first intermediate feature; The first intermediate feature is input into the decoder of the teacher model to obtain the first reconstructed image; A first loss function is obtained based on the first input image and the first reconstructed image; The teacher model is pre-trained based on the first loss function.
3. The knowledge distillation method for a mask autoencoder according to claim 1, characterized in that, The fine-tuning training of the pre-trained student model based on downstream tasks includes: The fourth intermediate feature of the training set for the downstream task is obtained from the encoder of the pre-trained student model. The fourth intermediate feature is input into the mapping layer of the downstream task for fine-tuning training.
4. A knowledge distillation device for a mask autoencoder, characterized in that, include: The model building module is used to build a teacher model and a student model of the mask autoencoder, respectively. The teacher model and the student model are both visual transformation models, and the size of the teacher model is larger than that of the student model. The teacher model training module is used for pre-training the teacher model; The knowledge distillation pre-training module is used to perform knowledge distillation pre-training on the student model based on the pre-trained teacher model. The fine-tuning training module is used to fine-tune the pre-trained student model based on downstream tasks; The knowledge distillation pre-training module includes: a first acquisition module, used to acquire reconstruction loss and distillation loss based on the student model and the pre-trained teacher model; a second loss module, used to determine a second loss function based on the reconstruction loss and distillation loss; and a second pre-training module, used to perform knowledge distillation pre-training on the student model based on the second loss function. The first acquisition module includes: a second masking module, used to mask the second input image according to a second preset masking rate; a second reconstruction module, used to input the masked second input image into the encoder of a student model to obtain second intermediate features, and input the second intermediate features into the decoder of a student model to obtain a second reconstructed image; a third reconstruction module, used to input the masked second input image into the encoder of a pre-trained teacher model to obtain third intermediate features, and input the third intermediate features into the decoder of a pre-trained teacher model to obtain a third reconstructed image; a second loss module, used to obtain a reconstruction loss based on the second reconstructed image; and a third loss module, used to obtain a distillation loss based on the second intermediate features, the third intermediate features, the second reconstructed image, and the third reconstructed image. The second loss module includes: a second acquisition module, used to acquire the distance between pixels in the masked region of the second input image and the second reconstructed image, and the distance between pixels in the unmasked region of the second input image and the second reconstructed image; and a reconstruction loss module, used to acquire a reconstruction loss based on the distance between pixels in the masked region and the distance between pixels in the unmasked region, wherein the calculation formula for the reconstruction loss is: In the formula, Indicates the reconstruction loss. , These represent the number of pixels in the masked area and the number of pixels in the unmasked area, respectively. The value of the i-th pixel in the second input image. This represents the value of the i-th pixel in the second reconstructed image. The third loss module includes: a calculation module for calculating the distance between the second intermediate feature and the third intermediate feature; a distance module for calculating the distance between pixels in the masked region of the second reconstructed image and the third reconstructed image based on a preset temperature coefficient; and a distillation loss module for obtaining distillation loss based on the distance between the second intermediate feature and the third intermediate feature, and the distance between pixels in the masked region of the second reconstructed image and the third reconstructed image. The formula for calculating the distillation loss is as follows: In the formula, The function is used to calculate the second intermediate feature. and the third intermediate feature The distance between them is determined by using Paradigm The third reconstructed image for the teacher model The value of the i-th pixel This is the preset temperature coefficient.
5. The knowledge distillation apparatus for a mask autoencoder according to claim 4, characterized in that, The teacher model training module includes: The first masking module is used to mask the first input image according to a first preset masking rate; The first input module is used to input the masked first input image into the encoder of the teacher model to obtain the first intermediate feature; The first reconstruction module is used to input the first intermediate features into the decoder of the teacher model to obtain the first reconstructed image; The first loss module is used to obtain a first loss function based on the first input image and the first reconstructed image; The first pre-training module is used to pre-train the teacher model based on the first loss function.
6. The knowledge distillation apparatus for a mask autoencoder according to claim 4, characterized in that, The fine-tuning training module includes: The third acquisition module is used to acquire the fourth intermediate feature of the training set of the downstream task based on the encoder of the pre-trained student model. The fine-tuning module is used to input the fourth intermediate feature into the mapping layer of the downstream task for fine-tuning training.
7. An electronic device, characterized in that, include: A memory and a processor are communicatively connected, the memory storing computer instructions, and the processor executing the computer instructions to perform the knowledge distillation method of the mask autoencoder as described in any one of claims 1 to 3.
8. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores computer instructions for causing the computer to perform the knowledge distillation method of the mask autoencoder as described in any one of claims 1 to 3.
Citation Information
Patent Citations
Mobile terminal image segmentation model training method and device
CN114494275A
Image moire removing method and device suitable for intelligent terminal and storage medium
CN114596479A