Neural network knowledge distillation method based on learnable feature transformation

Abstract

The application provides a neural network knowledge distillation method based on a learnable feature transformation, and belongs to the technical field of computer vision.The intermediate features and output results of a student model and a teacher model are aligned, a complex feature transformation module does not need to be designed for different tasks, complex hyperparameters are not introduced, a cumbersome parameter adjustment step is avoided, the universality of knowledge distillation on multiple tasks can be improved, the cumbersome structure design is avoided while the knowledge distillation effect is improved, and the performance is improved on multiple computer vision tasks (such as picture classification, target detection, semantic segmentation and the like).

Description

TECHNICAL FIELD

[0001] The present application belongs to the technical field of computer vision, and relates to computer vision, neural network model compression, neural network knowledge distillation based on intermediate features and other deep learning technologies. BACKGROUND

[0002] In recent years, with the continuous development of deep learning technology, deep convolutional neural networks are widely used in computer vision tasks such as image classification, object detection, semantic segmentation, and have achieved better and better performance. Behind the better performance, the complexity of the deep convolutional neural network model is also increasing, and the demand for computing resources and storage resources is increasing, making it difficult to deploy on resource-constrained devices such as mobile devices and embedded platforms. To solve this problem, neural network model compression technology is needed.

[0003] Knowledge distillation is an important method in the current neural network model compression technology. This method uses a large-scale neural network as a teacher network and a small-scale neural network as a student network, and transfers the knowledge of the teacher network to the student network, thereby obtaining a neural network with low complexity, good performance, and easy deployment, achieving the purpose of model compression.

[0004] Currently, the mainstream knowledge distillation methods are divided into output response-based and intermediate feature-based knowledge distillation. The output response-based knowledge distillation method uses the prediction results of the tail layer of the teacher model as supervision information to guide the student model to imitate the behavior of the teacher model. The intermediate feature-based knowledge distillation method uses the features of the intermediate hidden layer of the teacher model as a supervision signal to guide the student model training. In practical applications, various knowledge distillation methods have been derived for different visual tasks, and these methods often have many hand-designed parts such as loss functions and feature masks. These hand-designed parts reduce the generality of the distillation method and increase the difficulty of parameter tuning. SUMMARY

[0005] To solve the above problems, the present application provides a knowledge distillation method based on learnable feature transformation, which aligns the intermediate features and output responses of the student model with the teacher model, improves the knowledge distillation effect, and eliminates the complexity of hand-designed structures, achieving performance improvement in multiple computer vision tasks such as image classification, object detection, and semantic segmentation.

[0006] The technical solution provided by the present application is:

[0007] A knowledge distillation method based on learnable feature transformation, as shown in Figure 1 The steps include:

[0008] 1) inputting input data into a teacher model, the intermediate layer of the teacher model outputting a first feature map, and inputting the input data into a student model, the intermediate layer of the student model outputting a second feature map;

[0009] 2) aligning the second feature map with the first feature map in the spatial dimension and the channel dimension, the aligned feature map passing through a multi-layer perception module to obtain a third feature map; meanwhile, unfolding and transposing the shape of the aligned feature map, and then passing through another multi-layer perception module to obtain a transformed feature map, and then restoring the shape of the transformed feature map to the shape before transformation to obtain a fourth feature map;

[0010] 3) calculating the mean square error loss between the first feature map and the third feature map as a spatial feature loss, and calculating the mean square error loss between the first feature map and the fourth feature map as a channel feature loss, and weighting and summing the spatial feature loss and the channel feature loss as a knowledge distillation loss function between the teacher model and the student model;

[0011] 4) training the student model according to the knowledge distillation loss function to realize knowledge distillation.

[0012] Preferably, the multi-layer perception module is a multi-layer perception structure with one hidden layer and a ReLU activation function.

[0013] Preferably, the second feature map is aligned with the first feature map in the spatial dimension and the channel dimension through bilinear interpolation and 1x1 convolution.

[0014] Further, a downstream task of the student model is obtained, a target function of the model is matched according to the type of the downstream task, the target function and the knowledge distillation loss function are combined to train the student model.

[0015] Further, the hyperparameters of the distillation loss function are adjusted according to the teacher model, the student model and the downstream task, the regression loss function, the classification loss function and the knowledge distillation loss function in the target function are summed to obtain a total loss function for training the student model, and the student model is trained according to the total loss function.

[0016] The present application has the following beneficial effects:

[0017] The present application provides a knowledge distillation method based on learnable feature transformation, which aligns the features of the teacher model and the student model, improves the distillation effect, and at the same time, without designing a complex feature transformation module for different tasks, without introducing complex hyperparameters, without going through the cumbersome parameter adjustment step, and improving the universality of knowledge distillation on multiple tasks, good results can be achieved on various computer vision tasks. BRIEF DESCRIPTION OF DRAWINGS

[0018] Figure 1 The flowchart of the knowledge distillation method based on the learnable feature transformation of the present application is shown in the figure.

[0019] Figure 2 The architecture diagram of the training process of the student model of the embodiment of the present application is shown in the figure. DETAILED DESCRIPTION

[0020] The present application will be further described by examples in conjunction with the accompanying drawings, but the scope of the present application is not limited in any way.

[0021] Taking the large-scale target detection dataset COCO as an example, RetinNet-rx101 pre-trained on the dataset is taken as the teacher model, and RetinaNet-R50 is selected as the student model to illustrate how to perform knowledge distillation on the target detection task through the learnable transformation module, as shown in the figure. Figure 2

[0022] Step S1: inputting the input data into the teacher model to obtain the first feature map of the intermediate layer output of the teacher model, and inputting the input data into the student model to obtain the second feature map of the intermediate layer output of the student model, specifically comprising:

[0023] S11: inputting any batch of original training pictures into the teacher model RetinNet-rx101 to obtain the first feature map of the intermediate layer output in the FPN part of the teacher model.

[0024] S12: inputting the training pictures into the student model RetinaNet-R50 to obtain the second feature map of the intermediate layer output in the FPN part of the student model.

[0025] Step S2: obtaining the third feature map and the fourth feature map by using the multilayer perception module, specifically comprising:

[0026] S21: aligning the second feature map and the first feature map in the spatial dimension and the channel dimension by bilinear interpolation and 1x1 convolution to obtain the aligned feature map;

[0027] S22: obtaining the third feature map by passing the aligned feature map through a multilayer perception module with one hidden layer and ReLU as the activation function.

[0028] ​S23: assuming that the aligned feature map shape is [N, C, H, W], the shape of the feature map is adjusted to [N, (H*W), C] through unfolding and transposition operations, the adjusted feature map is obtained through a multilayer perception module with a hidden layer number of 1 and a ReLU activation function to obtain a transformed feature map, and the shape of the transformed feature map is adjusted to [N, C, H, W] to obtain the fourth feature map.

[0029] Step S3: according to the first feature map, the third feature map and the fourth feature map, the spatial feature loss and the channel feature loss between the teacher model and the student model are calculated, and the spatial feature loss and the channel feature loss are weighted and summed as the knowledge distillation loss function between the teacher model and the student model, specifically including:

[0030] S31: the mean square error loss between the first feature map and the third feature map is calculated as the spatial feature loss, and the expression is:

[0031]

[0032] wherein feat T is the first feature map, is the third feature map

[0033] S32: the mean square error loss between the first feature map and the fourth feature map is calculated as the channel feature loss, and the expression is:

[0034]

[0035] wherein feat T is the first feature map, is the fourth feature map

[0036] S33: the spatial feature loss and the channel feature loss are weighted and summed to obtain the knowledge distillation loss function, and the expression is:

[0037] L distill = aLoss Spatial + bLoss Channel

[0038] wherein a, b are hyperparameters, and in this embodiment, are set to 2e-5 and 1e-6 respectively.

[0039] Step S4: according to the knowledge distillation loss function, the student model is trained to realize knowledge distillation.

[0040] Further, the downstream task of the student model is obtained, and in this embodiment, the downstream task is a target detection task.

[0041] Step S5: According to the downstream task type matching model objective function, in this embodiment, the objective function of the model is divided into a regression loss function and a classification loss function, the expression of the regression loss function is:

[0042]

[0043] Where t i is the deviation of each predicted anchor from the Ground Truth (GT), and is the true deviation of each anchor from the GT.

[0044] In this embodiment, the classification loss function adopts Focal Loss, and the expression is:

[0045] L cls = -α t (1-p t )γlog(p t )

[0046] Where p t is the probability value of the sample being correctly classified, α t , and γ are hyperparameters, which are set to 0.25 and 2.0 respectively in this embodiment.

[0047] Step S6: Adjust the hyperparameters of the distillation loss function according to the teacher model, the student model, and the downstream task, obtain the total loss function of the student model training by the objective function, the knowledge distillation loss function, and the hyperparameters; train the student model according to the total loss function, and the expression of the total loss function is:

[0048] L total = L reg + L cls + L distill .

[0049] For the image classification task, the results on the ImageNet dataset show that using ResNet34 as the teacher model, ResNet18 as the student model, and the knowledge distillation method proposed in the application, the Top-1 accuracy on the test set can be improved from 69.9% to 71.4%; for the target detection task, the results on the MSCOCO dataset show that using RetinaNet-RX101 as the teacher model, RetinaNet-R50 as the student model, and the knowledge distillation method proposed in the application, the mAP of the student model can be improved from 37.4% to 41.0%; for the semantic segmentation task, the results on the CityScapes dataset show that using PSPNet-ResNet34 as the teacher model, PSPNet-ResNet18 as the student model, and the knowledge distillation method proposed in the application, the mIoU of the student model can be improved from 69.9% to 74.2% (Note: ImageNet is a large-scale image classification dataset, Top1-accuracy is used to measure the image classification accuracy; MSCOCO is a large-scale dataset containing target detection tasks, and the mAP of bbox is an indicator for measuring the performance of target detection; CityScapes is a semantic segmentation dataset, and mIoU is an indicator for measuring the performance of semantic segmentation). In addition, the application can also be used to realize cross-model knowledge distillation and achieve good results. For example, for the image classification task, on the Cifar-100 dataset, using ResNet56 based on the convolutional neural network architecture as the teacher model and ViT-tiny based on the Transformer architecture as the student model, the Top1-accuracy of the student model can be improved from 57.8% to 77.5% (Note: Cifar100 is a small-scale image classification dataset).

[0050] The application is described above through detailed implementation cases, and researchers and technicians in the field can make non-substantial changes in form or content according to the above steps without deviating from the scope of the essential protection of the application. Therefore, the application is not limited to the content disclosed in the above examples, and the protection scope of the application should be subject to the description of the claims.

Claims

1. A method of knowledge distillation based on learnable feature transformation, characterized in that, The steps include: 1) inputting input data into a teacher model, the intermediate layer of the teacher model outputting a first feature map, and inputting the input data into a student model, the intermediate layer of the student model outputting a second feature map; 2) aligning the second feature map with the first feature map in the spatial dimension and the channel dimension, the aligned feature map obtaining a third feature map through a multi-layer perception module; at the same time, unfolding and transposing the shape of the aligned feature map, and then obtaining a transformed feature map through another multi-layer perception module, and then restoring the shape of the transformed feature map to the shape before transformation to obtain a fourth feature map; specifically including: S21: aligning the second feature map with the first feature map in the spatial dimension and the channel dimension through bilinear interpolation and 1x1 convolution to obtain an aligned feature map; S22: obtaining a third feature map by passing the aligned feature map through a multi-layer perception module with one hidden layer and a ReLU activation function; S23: assuming that the shape of the aligned feature map is [N, C, H, W], adjusting the shape of the feature map to [N, (H*W), C] through unfolding and transposing operations, and obtaining a transformed feature map by passing the adjusted feature map through a multi-layer perception module with one hidden layer and a ReLU activation function, and then adjusting the shape of the transformed feature map to [N, C, H, W] to obtain the fourth feature map; 3) calculating the mean square error loss between the first feature map and the third feature map as a spatial feature loss, calculating the mean square error loss between the first feature map and the fourth feature map as a channel feature loss, and weighting and summing the spatial feature loss and the channel feature loss as a knowledge distillation loss function between the teacher model and the student model; specifically including: S31: calculating the mean square error loss between the first feature map and the third feature map as the spatial feature loss, the expression of which is: wherein is the first feature map, is the third feature map; S32: calculating the mean square error loss between the first feature map and the fourth feature map as the channel feature loss, the expression of which is: ； wherein is the first feature map, is the fourth feature map S33: weighting and summing the spatial feature loss and the channel feature loss to obtain the knowledge distillation loss function, the expression of which is: ； 4) training the student model according to the knowledge distillation loss function to realize knowledge distillation.

2. The knowledge distillation method based on learnable feature transformation of claim 1, wherein, In step 4), the downstream task of the student model is obtained, the target function of the model is matched according to the type of the downstream task, the target function and the knowledge distillation loss function are combined, and the student model is trained.

3. The knowledge distillation method based on learnable feature transformation of claim 2, wherein, In step 4), the regression loss function, the classification loss function and the knowledge distillation loss function in the target function are summed to obtain a total loss function for training the student model, and the student model is trained according to the total loss function.

Images