A pre-training method based on MAE
By performing two-stage masking and downsampling using the improved MAE algorithm, the problem of high GPU memory requirements in vision pre-training methods is solved, enabling efficient training of the Swin model on general-purpose machines and supporting fine-tuning of supervised data and downstream tasks.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- NEWLAND DIGITAL TECH CO LTD
- Filing Date
- 2022-12-29
- Publication Date
- 2026-07-14
AI Technical Summary
Existing pre-training methods in the field of vision require a large amount of GPU memory resources, making it difficult to train on general-purpose machines. In particular, the SimMIM algorithm of the Swin series models has low computational efficiency and consumes a huge amount of GPU memory, making it impossible to effectively utilize the resources of general-purpose machines.
An improved MAE algorithm is used for two-stage masking. First, patch grid region segmentation and random sampling are performed on the training set images. Then, a token random masking operation is performed. Downsampling is performed in combination with the LinearEmbedding and transformerblock modules, which is suitable for the pyramid structure of the swin pre-trained model.
The memory requirement has been reduced from over 100GB to 20GB, enabling training on dual-card machines, improving the utilization efficiency of computing resources, and supporting the fine-tuning of the swin model with supervised data in various scenarios and downstream tasks.
Abstract
Description
Technical Field
[0001] This invention is applied to the field of computer vision, specifically a pre-training method based on MAE. Background Technology
[0002] Pre-training is a mainstream approach in computer vision. Originating in computer vision, pre-training the backbone of object detection and classification models on ImageNet to assist in other data scenarios. This pre-training method is based on the CNN architecture and was later adopted in natural language processing, achieving great success with the rise of the transformer model architecture. Currently, cutting-edge research in computer vision focuses on applying the transformer-based pre-trained models that have been highly successful in natural language processing to the field of computer vision.
[0003] There are two main types of unsupervised pre-training methods in the field of computer vision: one is contrastive learning, which is suitable for classification tasks. Through data augmentation, the augmented image is considered a positive sample, while the remaining images are considered negative samples. The other is learning from natural language processing pre-training methods, which is based on the prediction of masks for unsupervised learning. Currently, there are two main masking methods in the field of computer vision: one is the MAE method represented by the VIT pre-trained model, and the other is the SimMIM method represented by the Swin pre-trained model.
[0004] However, since pre-training methods in the field of vision require very large amounts of GPU memory, the aforementioned methods often need to be trained on a server with 8 GPUs or more. Taking the ImageNet dataset as an example, it takes about 15 days to run on a single 8-GPU Titan RTX server, and a batch of data requires more than 120GB of GPU memory. Therefore, pre-training methods have high requirements for server resources, especially GPU memory. When the production environment does not have large server resources, pre-training tasks become difficult to perform.
[0005] The MAE algorithm reconstructs the entire image. For the entire image, it randomly samples 25% of the pixels to predict the remaining 75% of the masked pixels. It is computationally efficient, but this algorithm only supports pre-trained models based on the VIT structure because the VIT model structure supports global feature extraction. The SimMIM algorithm first divides the image into patch regions, then randomly masks all pixels in some patch regions, and uses the remaining pixels to predict these masked pixels. It supports pre-trained models with pyramid structures like the Swin model. It divides the image region by continuously downsampling and then masks some regions. It must receive a full-size image as input, so its computational efficiency is relatively low.
[0006] Currently, the SWIN pre-trained model performs better than VIT on detection tasks. Therefore, the SWIN series of models are the first choice when performing detection tasks. However, the SWIN series of models supports the SimMIM algorithm for pre-training, which has low computational efficiency and consumes a lot of GPU memory.
[0007] Therefore, how to improve the algorithm to combine the advantages of MAE (which does not require full image input and has high computational efficiency) and SimMIM (which can extract pyramid structure pixels) and enable pre-training tasks on the Swin series models has become a current technical challenge. Summary of the Invention
[0008] The technical problem to be solved by the present invention is to provide a pre-training method based on MAE to address the shortcomings of the existing technology.
[0009] To address the aforementioned technical problems, the present invention provides a MAE-based pre-training method, which specifically includes the following steps:
[0010] Obtain training set images;
[0011] The improved MAE algorithm is used to perform two-stage masking processing on the training set images.
[0012] The masked image is input into the pre-trained model to obtain features;
[0013] The features are processed by the pre-trained model to finally output the image with the completed mask.
[0014] As one possible implementation, the two-stage masking process for the training set images using the improved MAE algorithm specifically includes:
[0015] In the first stage, the training set images are segmented into patch grid regions. One grid is randomly sampled from every 4 grids in a 2x2 grid to ensure that the operators are evenly distributed to an equal number of token sequence objects.
[0016] In the second stage, based on the token objects sampled in the first stage, another random masking operation is performed, randomly sampling a preset proportion of the tokens as learnable shared mask tokens.
[0017] As one possible implementation, further, the predetermined proportion of the tokens randomly sampled as a predetermined proportion in a learnable shared mask token is 25%.
[0018] As one possible implementation, the step of inputting the masked image into the pre-trained model to obtain features specifically includes:
[0019] Input the masked image into the Swin model;
[0020] Downsampling is performed by 4 times using two LinearEmbedding + TransformerBlock modules;
[0021] Downsampling was performed by 8 times using two LinearEmbedding + TransformerBlock modules;
[0022] Downsampling was performed by 16 times using two LinearEmbedding + TransformerBlock modules;
[0023] Features are obtained through feature calculation using two sliding window W-MSA+SW-MSA modules.
[0024] As one possible implementation, the step of the feature being processed by the pre-trained model to finally output the image with the padded mask is as follows: input the feature into the Swin model and output the image with the padded mask.
[0025] As one possible implementation, the step of inputting the masked image into the pre-trained model to obtain features specifically includes:
[0026] Input the masked image into the Vit model;
[0027] The masked image features are transformed into a sequential sequence using six transformer encoder structures.
[0028] The features are then decoded into a feature sequence of the completed image using six transformer decoder structures to obtain the features.
[0029] As one possible implementation, the step of the feature being processed by the pre-trained model to finally output the image with the padded mask is as follows: input the feature into the Vit model and output the image with the padded mask.
[0030] The present invention adopts the above technical solution and has the following beneficial effects:
[0031] This invention proposes an improvement to the MAE unsupervised pre-training algorithm. The improved MAE algorithm can perform unsupervised training on a SWIN pre-trained model that supports a pyramid structure, reducing the required GPU memory from over 100GB to 20GB. This allows training to be supported on dual-GPU machines, effectively saving a significant amount of GPU memory and enabling large-scale pre-training tasks on general machine resources, thus significantly improving the utilization of computing resources. The SWIN model pre-trained using MAE unsupervised methods can support fine-tuning of supervised data in various scenarios and training of downstream tasks. Detailed Implementation
[0032] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions in the embodiments of the present invention will be clearly and completely described below.
[0033] Example 1:
[0034] This invention provides a pre-training method based on MAE, specifically including the following steps:
[0035] Obtain training set images;
[0036] The improved MAE algorithm is used to perform two-stage masking processing on the training set images.
[0037] The masked image is input into the pre-trained model to obtain features;
[0038] The features are processed by the pre-trained model to finally output the image with the completed mask.
[0039] As one possible implementation, the two-stage masking process for the training set images using the improved MAE algorithm specifically includes:
[0040] In the first stage, the training set images are segmented into patch grid regions. One grid cell is randomly sampled from every 2x2 four grid cells to ensure that operators are evenly distributed to an equal number of token sequence objects. This operation ensures that all global window operators in the pyramid ViT can be evenly distributed to an equal number of token sequence objects.
[0041] In the second stage, based on the tokens sampled in the first stage, a second random masking operation is performed, randomly sampling a predetermined proportion of the tokens into learnable shared mask tokens. Essentially, this operation removes some information because the image patch blocks sampled in the first step are spatially more uniform than the random sampling in MAE. This uniformity reduces the difficulty of the pre-training task, as the network can easily infer the approximate pixels of the missing parts using surrounding low-level pixel information, thus weakening the network's ability to learn high-level semantics. The second stage breaks the uniformity of the first stage, allowing the network to learn higher-quality semantic representations through self-supervised tasks, while modeling as learnable mask tokens ensures the equivalence of local window operator operation objects. The predetermined proportion of tokens randomly sampled into learnable shared mask tokens is 25%. After a two-stage masking process, the original HxW input image can be reconstructed into a new image of size (H / 2)x(W / 2). Inputting this new image into the pyramid structure swin reduces complexity, achieves an efficient input mode for MAE, and is also applicable to pyramid structure models.
[0042] Taking the application to the Swin pre-trained model as an example: the step of inputting the masked image into the pre-trained model to obtain features specifically includes:
[0043] Input the masked image into the Swin model;
[0044] Downsampling is performed by 4 times using two LinearEmbedding + TransformerBlock modules;
[0045] Downsampling was performed by 8 times using two LinearEmbedding + TransformerBlock modules;
[0046] Downsampling was performed by 16 times using two LinearEmbedding + TransformerBlock modules;
[0047] Features are obtained through feature calculation using two sliding window W-MSA+SW-MSA modules.
[0048] The specific steps for the feature to be processed by the pre-trained model and finally output as a padded masked image are as follows: input the features into the Swin model and output the padded masked image.
[0049] Taking the application to the Vit pre-trained model as an example: the step of inputting the masked image into the pre-trained model to obtain features specifically includes:
[0050] Input the masked image into the Vit model;
[0051] The masked image features are transformed into a sequential sequence using six transformer encoder structures.
[0052] The features are then decoded into a feature sequence of the completed image using six transformer decoder structures to obtain the features.
[0053] The specific steps for the feature to be processed by the pre-trained model and finally output as a padded masked image are as follows: input the feature into the Vit model and output the padded masked image.
[0054] Example 2:
[0055] The main difference between this embodiment and Embodiment 1 is that the improved MAE algorithm is used to perform unsupervised training of the swin pre-trained model on an unsupervised dataset of a fireworks detection scenario, with the training batch size set to 72.
[0056] This embodiment only requires 20GB of VRAM during training, and two 3060 graphics cards (12GB each) are sufficient to support training. Compared to SimMIM and MAE, which require more than 100GB of VRAM, this represents a significant optimization in VRAM requirements. The unsupervised dataset size is 53,072 images, and the training iterations totaled 200 epochs, with a final training time of 6 days and 2 hours. The Swin model, after unsupervised pre-training, can continue to be fine-tuned on supervised data in fireworks detection scenarios and trained for downstream detection tasks.
[0057] The above description is an embodiment of the present invention. For those skilled in the art, any equivalent changes, modifications, substitutions and variations made in accordance with the teachings of the present invention without departing from the principles and spirit of the present invention should be covered by the present invention.
Claims
1. A pre-training method based on MAE, characterized in that, Specifically, the steps include the following: Obtain training set images; The improved MAE algorithm is used to perform two-stage masking processing on the training set images. The masked image is input into the pre-trained model to obtain features; The features are processed by the pre-trained model, which ultimately outputs the image with the completed mask. The two-stage masking process for the training set images using the improved MAE algorithm specifically includes: In the first stage, the training set images are segmented into patch grid regions. One grid is randomly sampled from every 4 grids in a 2x2 grid to ensure that the operators are evenly distributed to an equal number of token sequence objects. In the second stage, based on the token objects sampled in the first stage, another random masking operation is performed, randomly sampling a preset proportion of the tokens as learnable shared mask tokens.
2. The MAE-based pre-training method according to claim 1, characterized in that: The preset proportion of the tokens in the process of randomly sampling a preset proportion of the tokens into a learnable shared mask token is 25%.
3. The MAE-based pre-training method according to claim 1, characterized in that: The step of inputting the masked image into the pre-trained model to obtain features specifically includes: Input the masked image into the Swin model; Downsampling is achieved by a factor of 4 using two Linear Embedding + Transformer block modules; Downsampling is achieved by 8 times using two Linear Embedding + Transformer block modules; Downsampling is achieved by 16 times using two Linear Embedding + Transformer block modules; Features are obtained through feature calculation using two sliding window W-MSA+SW-MSA modules.
4. The MAE-based pre-training method according to claim 3, characterized in that: The specific steps for the features to be processed by the pre-trained model and finally output as a padded masked image are as follows: input the features into the Swin model and output the padded masked image.
5. The MAE-based pre-training method according to claim 1, characterized in that: The step of inputting the masked image into the pre-trained model to obtain features specifically includes: Input the masked image into the Vit model; The masked image features are transformed into a sequential sequence using a 6-transformer encoder structure. The features are then decoded into a feature sequence of the completed image using six transformer decoder structures to obtain the features.
6. The MAE-based pre-training method according to claim 5, characterized in that: The specific steps for the feature to be processed by the pre-trained model and finally output as a padded masked image are as follows: input the feature into the Vit model and output the padded masked image.