A multi-stage training method of a face key point detection model based on a ViT network
By employing a multi-stage training method, combining pseudo-labeled datasets and manually labeled datasets, the face landmark detection model of the ViT network is optimized, solving the problems of poor robustness and high data volume under occlusion conditions, and achieving efficient landmark detection and cross-scene adaptability.
Patent Information
- Application Number
- CN202211086101.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2022-09-06
- Publication Date
- 2026-02-06
- Estimated Expiration
- 2042-09-06
AI Technical Summary
Existing facial landmark detection methods based on ViT networks have poor robustness to occlusion and require a large amount of data, making them difficult to effectively transfer to different scenarios.
A multi-stage training method is adopted. First, dual-branch training is performed on a pseudo-label dataset, including an unsupervised training branch for MAE and a supervised training branch for face key point detection. Initial training is performed on the pseudo-label dataset, followed by supervised training on a manually labeled dataset. The method combines fully connected layers, global pooling layers, and multilayer perceptrons, and uses the AdamW optimizer and cosine annealing learning rate strategy.
It improves the robustness of facial landmark detection, especially with excellent performance under occlusion conditions, reduces the requirement for the amount of labeled data, and makes the model easier to transfer to different scenarios.
Smart Images
Figure CN115311728B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The application relates to the technical field of computer vision and artificial intelligence, and particularly relates to a multi-stage training method of a face key point detection model based on a ViT network. BACKGROUND
[0002] Face key point detection is an important branch of computer vision and is used for positioning various nodes of a face, such as eyes, a nose and a mouth, and the state of the target face can be further understood according to the detected key points. Most of the current face key point detection methods are based on convolutional neural networks, and the key point detection accuracy of the methods is greatly affected by occlusions such as wearing a mask and wearing glasses. Recently, the ViT (Vision Transformer) network composed of a Transformer structure has more and more obvious advantages in the field of computer vision, and the global attention different from the convolutional network makes the ViT network have good performance in different scenes and strong robustness, so the face key point detection network is built based on the ViT architecture. In view of the problems of the ViT network architecture, such as large data requirement and difficulty in convergence, a multi-stage training method of a face key point detection model based on a ViT network is provided, which reduces the requirement of the ViT model for the amount of labeled data and guarantees the training effect of the model. SUMMARY
[0003] The application aims to provide a multi-stage training method of a face key point detection model based on a ViT network to solve the problems in the background.
[0004] To achieve the above object, the application provides the following technical scheme: a multi-stage training method of a face key point detection model based on a ViT network, comprising the following steps:
[0005] Step 1, constructing a double-branch training structure based on a ViT network;
[0006] S1, for the MAE unsupervised training branch, adding a full connection layer and a decoder network composed of Transformer blocks after the ViT network, and the decoder network comprises 8 Transformer blocks, and for the face key point training branch, adding a global pooling layer and a multilayer perceptron after the ViT network;
[0007] S2, training the double-branch training structure by using a pseudo-label data set in the first stage, wherein the pseudo-label data set is constructed by predicting a face picture by using a PFLD face key point detection model;
[0008] S3, as an example of a single iteration process of training, the picture needs to go through image block projection PatchEmbedding and position encoding before inputting into the ViT network. Image block projection is to divide the input image into sub-patches;
[0009] S4, the output feature is transmitted into the MAE unsupervised training branch. First, the input is changed in feature dimension through a fully connected layer. The changed feature dimension is D FC1 , then the fully connected layer outputs the feature BxKxD FC1 ; the output feature is filled into Bx((H / h)*(W / w))xD FC1 using a learnable tensor Token according to the adopted index Indices, and is sent into the decoding network after reposition encoding. The feature output by the decoding network is restored to the picture tensor BxHxWx3 through a fully connected layer and is calculated with the input image in pixel level MSE loss. The loss calculation is only performed on the image blocks not adopted;
[0010] S5, the input feature is transmitted into the face key point detection supervised training branch. The feature goes through a global pooling layer and a multi-layer perceptron MLP, which contains two fully connected layers. The first fully connected layer is the hidden layer, and the second fully connected layer is the prediction layer. The output tensor of the prediction layer is Bxlandmarks*2, where landmarks is the number of predicted face key points. The output tensor is calculated with the label in Wing loss;
[0011] Step 2, only the network weights of the supervised training branch are used for face key point detection supervised fine-tune training. The network is fine-tuned for face key point detection using a manually labeled face key point detection dataset. The training adopts AdamW optimizer, batch size is 512, learning rate is 1e -4 , cosine annealing learning rate reduction strategy is adopted, and data augmentation strategies include random graying and random color disturbance;
[0012] After training on the manually labeled dataset until the model converges, the trained face key point model is obtained after iterative training.
[0013] Preferably, in S1, the multi-layer perceptron contains two fully connected layers. The first fully connected layer is the hidden layer, and the second fully connected layer is the key point prediction layer.
[0014] Preferably, in S2, due to the prediction error of the model, the pseudo-label dataset is only used in the first stage.
[0015] Preferably, in S3, ViT-S is adopted, and the feature size output by ViT is BxKx768.
[0016] Preferably, in S5, the total loss function is Loss = a * Loss MSE + b * Loss Wing .
[0017] Preferably, in step 2, at the same time, in order to avoid overfitting of the model in the training process, since the manually annotated data set is usually small, R is set to 0.9 during training, that is, 10% of each picture is randomly discarded during training.
[0018] The multi-stage training method of the face key point detection model based on the ViT network has the beneficial effects that:
[0019] 1. The face key point detection model based on the ViT network has good detection effect on face key points in different scenes, especially in occlusion state.
[0020] 2. The multi-stage training method reduces the data volume requirement of the model based on the ViT network for labeled data, and is easy to migrate to different scenes. BRIEF DESCRIPTION OF DRAWINGS
[0021] Fig. 1 is a double-branch training schematic diagram of the present application;
[0022] Fig. 2 is a face key point supervised training schematic diagram of the present application. DETAILED DESCRIPTION
[0023] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor are within the scope of protection of the present application.
[0024] Embodiments, please refer to Figs. 1-2 The present application provides a technical solution: a multi-stage training method of a face key point detection model based on a ViT network, comprising the following steps:
[0025] Step 1: constructing a double-branch training structure based on the ViT network, including a MAE (Masked Autoencoders) unsupervised training branch and a face key point detection supervised training branch; using a pseudo-label training data set to train the constructed training structure in one stage;
[0026] S1, for the MAE unsupervised training branch, add a fully connected layer after the ViT network and a decoder network composed of Transformer blocks, the decoder network contains 8 Transformer blocks, for the human face key point training branch, add a global pooling layer after the ViT network and a multilayer perceptron, the multilayer perceptron contains two fully connected layers, the first fully connected layer is a hidden layer, and the second fully connected layer is a key point prediction layer;
[0027] S2, the first stage of training, the double-branch training structure is trained using a pseudo-label dataset, which is constructed by predicting human face pictures using the PFLD human face key point detection model. Due to the prediction error of the model, this pseudo-label dataset is only used in the first stage;
[0028] S3, taking a single iteration process as an example, the picture needs to go through image block projection PatchEmbedding and position encoding before inputting into the ViT network. Image block projection is to divide the input image into sub-patches. Assuming that the input image resolution is HxW and the sub-patch size is w x h, the number of sub-patches is H / h x W / w, and the divided sub-patches are transformed into new 2-dimensional vectors and then expanded into 1-dimensional vectors, followed by embedding 1-dimensional explicit sin-cos position encoding information. The output vector dimension after the above operation is B x ((H / h)*(W / w)) x (w*h*3), where B is the input picture batch size BatchSize; for the output vector, sample indices Indices are generated according to the sampling ratio R by uniform distribution sampling, and K image blocks are selected by the sampling indices, K is the product of the sampling ratio R and the total number of image blocks, then the feature size of the Transformer module in ViT is B x K x (w*h*3), and ViT-S is used, the feature size output by ViT is B x K x 768;
[0029] S4, the output feature is transmitted into the MAE unsupervised training branch. First, the input is changed in feature dimension by a fully connected layer. Assuming that the changed feature dimension is D FC1 , the fully connected layer output feature is B x K x D FC1 ; the output feature is filled into B x ((H / h)*(W / w)) x D FC1 using a learnable tensor Token according to the sampling index Indices, and then sent into the decoding network after repositioning. The feature output by the decoding network is again passed through a fully connected layer and restored to a picture tensor of B x H x W x 3, and the pixel-level MSE loss is calculated with the input image, where the loss calculation is only performed on the image blocks that are not used;
[0030] S5, the input feature is transmitted into the face key point detection supervision training branch, the feature passes through a global pooling layer and a multilayer perceptron (MLP), wherein the multilayer perceptron includes two fully connected layers, the first fully connected layer is a hidden layer, and the second fully connected layer is a prediction layer, and the output tensor of the prediction layer is B x landmarks * 2, wherein landmarks is the number of predicted face key points, and the output tensor is subjected to Wing loss calculation with a label;
[0031] The total loss function is Loss = a * Loss MSE + β * Loss Wing ;
[0032] In the first stage training setting, a is 1, β is 0.01, the input resolution of the image is 112 x 112, the sub-image size is 8 x 8, and the feature B x 196 x 768 is obtained after image block projection and position coding; for 196 sub-image blocks of each image, the sampling ratio R is 0.25, K is 49, and the feature size B x 49 x 768 is sent into the ViT network Transformer part. FC1 D is 512, the number of Transformer blocks of the decoding network is 8, the feature dimension is 512, the feature dimension of the multilayer perceptron hidden layer of the face key point detection supervision training branch is 512, and landmarks is 98; the training adopts an AdamW optimizer, the batch size is 512, the learning rate is 1e -4 , a cosine annealing learning rate reduction strategy is adopted, and the data enhancement strategy only includes random cropping and random scaling. The model is trained on pseudo-label data for 400 epochs until convergence.
[0033] Step 2, two-stage training, remove the MAE unsupervised training branch in step 1, keep the face key point supervision training branch, load the pre-trained model weight and use the manually labeled face key point dataset for supervised training to obtain a face key point detection model.
[0034] The network weight of only keeping the supervision training branch is used for face key point detection supervision fine-tune training; the network is subjected to face key point supervision fine-tune training using a manually labeled face key point detection dataset, the training adopts an AdamW optimizer, the batch size is 512, the learning rate is 1e -4 , a cosine annealing learning rate reduction strategy is adopted, and the data enhancement strategy includes random graying and random color disturbance.
[0035] At the same time, since the manually labeled dataset is usually small, in order to avoid overfitting of the model during training, R is set to 0.9 during training, that is, 10% of each image is randomly discarded during training, and this method is more efficient and better than Dropout during training.
[0036] The face key point model is trained on the artificial labeled data set until the model converges, and the trained face key point model is obtained after iterative training.
[0037] The face key point detection model based on the ViT network has better robustness than the face key point detection algorithm based on the convolutional neural network, and still has good key point positioning ability for face occlusion. The multi-stage training method includes: stage one, double-branch training based on pseudo-label pre-training data set, the double-branch includes MAE (Masked Autoencoders) image reconstruction unsupervised training branch and face key point detection supervised training branch; stage two, face key point detection supervised training based on hand-labeled data set; the training method effectively reduces the data amount requirement of the ViT model for labeled data while ensuring the training effect of the model, so that the model is easy to migrate to different scenes.
[0038] Although embodiments of the present application have been shown and described, it is to be understood that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the present application, the scope of which is defined by the appended claims and their equivalents.
Claims
1. A multi-stage training method of a face key point detection model based on a ViT network, characterized in that, The method comprises the following steps: Step 1, constructing a double-branch training structure based on a ViT network; S1, for the MAE unsupervised training branch, adding a fully connected layer and a decoder network composed of Transformer blocks after the ViT network, the decoder network comprising 8 Transformer blocks, and for the human face key point training branch, adding a global pooling layer and a multilayer perceptron after the ViT network; S2, training of the first stage, training the double-branch training structure using a pseudo-label dataset, wherein the pseudo-label dataset is constructed by predicting human face pictures using a PFLD human face key point detection model; S3, taking a single iteration process of training as an example, the picture needs to undergo image block projection PatchEmbedding and position encoding before being input into the ViT network, the image block projection being to divide the input image into sub-pictures Patch; S4, the output feature is input into the MAE unsupervised training branch, first input through a fully connected layer to change the feature dimension, and the changed feature dimension is D FC1 , the output feature of the fully connected layer is BxKxD FC1 ; the output feature is filled into Bx((H / h)*(W / w))xD using a learnable tensor Token according to the adopted index Indices FC1 , and is sent into the decoding network after reposition encoding. The output feature of the decoding network is again processed through a fully connected layer and restored to a picture tensor of BxHxWx3, and a pixel-level MSE loss calculation is performed with the input image, wherein the loss calculation is only performed on the image blocks that are not adopted; S5, inputting the features into the human face key point detection supervised training branch, the features passing through a global pooling layer and a multilayer perceptron MLP, wherein the multilayer perceptron comprises two fully connected layers, the first fully connected layer being a hidden layer, and the second fully connected layer being a prediction layer, the prediction layer outputting a tensor of Bxlandmarks*2, wherein landmarks is the number of predicted human face key points, and the output tensor is subjected to Wing loss calculation with a label; Step 2, only the network weights of the supervised training branch are reserved for face key point detection supervised fine-tune training; the network is subjected to face key point supervised fine-tune training using a manually labeled face key point detection dataset, the training adopts an AdamW optimizer, a batch size of 512, and a learning rate of 1e -4 Cosine annealing learning rate reduction strategy, and data enhancement strategies include random graying and random color perturbation; Training on the manually annotated dataset until the model converges, and obtaining a trained human face key point model after iterative training.
2. The multi-stage training method of a ViT network-based face key point detection model according to claim 1, characterized in that: In S1, the multilayer perceptron comprises two fully connected layers, the first fully connected layer being a hidden layer, and the second fully connected layer being a key point prediction layer.
3. The multi-stage training method of a ViT network-based face key point detection model according to claim 1, characterized in that: In S2, since there is an error in model prediction, the pseudo-label dataset is only used in the first stage.
4. The multi-stage training method of a ViT network-based face key point detection model according to claim 1, characterized in that: In S3, ViT-S is used, and the feature size output by the ViT is BxKx768.
5. The multi-stage training method of a ViT network-based face key point detection model according to claim 1, characterized in that: In S5, the total loss function is Loss = a*Loss MSE + β*Loss Wing .
6. The multi-stage training method of a ViT network-based face key point detection model according to claim 1, characterized in that: In step 2, at the same time, since the manually annotated dataset is usually small, in order to avoid overfitting of the model during training, R is set to 0.9 during training, that is, 10% of each picture is randomly discarded during training.