Aerial image object detection method based on cross-level feature distillation

Abstract

The present invention discloses a method for aerial image target detection based on cross-level feature distillation, including: extracting a pre-trained teacher network mask map and a student network mask map, constructing a query module based on the pre-trained teacher network mask map and the student network mask map, performing cross-level feature distillation on targets of different scales based on the query module, using the multi-layer feature information of the pre-trained teacher network to guide the single-layer feature information of the student network, and enhancing the semantic information matching degree between the pre-trained teacher network and the student network. The present invention proposes a strategy of cross-level feature distillation based on the query module, which allows the student network to better learn the knowledge from the pre-trained teacher network by making a loss on the corresponding scale through the query module, and at the same time, can improve the detection performance of the student network without increasing the number of network parameters, and has certain advantages in the deployment of satellite platforms.

Description

Technical Field

[0001] The present invention relates to the field of aerial image application technology, and in particular to an aerial image target detection method based on cross-level feature distillation. Background Art

[0002] In the field of aerial image detection, if you want to achieve high-precision detection results, you generally need a larger deep neural network model, which consumes more computing resources and storage space. These large models are difficult to deploy on satellite platforms with high speed requirements.

[0003] In order to solve the above problems, some lightweight deep learning algorithms continue to emerge, such as knowledge distillation, network pruning, quantization, weight sharing, lightweight, etc. Some of these methods have achieved real-time requirements, but the detection accuracy still cannot meet the needs of the task. At present, in the field of target detection, knowledge distillation methods have gradually begun to be used to improve the performance of deep neural networks, balancing detection speed and detection accuracy. Knowledge distillation refers to transferring the feature extraction rules learned by the large model to the small model, improving the accuracy of the small model, and then using the small model in actual scenarios to achieve the purpose of model compression.

[0004] Knowledge distillation was first applied to the field of image classification, and it can transfer key knowledge such as network output, feature maps, and information flow. However, in the target detection task, there is a serious class imbalance problem, and the detection task is more difficult than the classification task. Directly applying the distillation method in image classification does not work well. Therefore, it is an important topic for researchers in this field to choose a suitable distillation method without increasing the amount of calculation. Summary of the invention

[0005] The purpose of the present invention is to use the knowledge distillation method to solve the problem that the existing aerial image detection algorithms are difficult to deploy on satellite platforms. The traditional knowledge distillation method uses a single teacher layer feature to distill a single student layer feature, which can easily lead to poor detection results due to semantic mismatch. In order to solve the problem more effectively, the present invention provides an aerial image target detection method based on cross-level feature distillation. The present invention first constructs a distillation basic model based on a pre-trained teacher network and a student network, and then constructs a teacher-student learning network framework based on the distillation basic model. Based on the query module, the appropriate area is selected for cross-level feature distillation. The single-layer feature information of the student network is guided by the multi-layer feature information of the pre-trained teacher network, thereby enhancing the generalization ability and robustness of the model.

[0006] The technical solution to achieve the purpose of the present invention is: a method for aerial image target detection based on cross-level feature distillation, comprising the following steps:

[0007] Step S1, obtain the aerial image detection dataset, randomly divide it into a training dataset and a test dataset, cut the images in the above two datasets into M×N blocks one by one, M and N represent the width and height of the image respectively, and the sizes of M and N are both between 800 and 1024. The training dataset is used to train the teacher network and the student network, and the test dataset is used to verify the student network after distillation, and then go to step S2.

[0008] Step S2: input the training data set into the teacher network for pre-training, obtain the pre-trained teacher network, and proceed to step S3.

[0009] Step S3: construct a distillation basic model based on the pre-trained teacher network and student network, and construct a teacher-student learning network framework based on cross-level feature distillation based on the distillation basic model, and then proceed to step S4.

[0010] Step S4, input the training data set into the teacher-student learning network framework, extract the features of each layer of the pre-trained teacher network and the teacher heat map of the query branch, and extract the features of each layer of the student network and the student heat map of the query branch, set a threshold to suppress the background pixels and low-confidence pixels in the teacher heat map and the student heat map, and obtain the corresponding pre-trained teacher network mask map and student network mask map, and construct a query loss function of the student network mask map and the target true value mask map; construct a query module based on the pre-trained teacher network mask map and the student network mask map, perform cross-level feature distillation on targets of different scales based on the query module, and use the multi-layer feature information of the pre-trained teacher network to guide the single-layer feature information of the student network; fix the parameters in the pre-trained teacher network, continuously update the training parameters of the student network by constructing the loss function of the student network, finally obtain the trained student network, and go to step S5.

[0011] Step S5: Use the trained student network to detect all the block images on the test data set, merge the detection results on the block images, and finally output the categories and regression positions of all targets in the test data set in the aerial images, and finally obtain the accuracy of the trained student network.

[0012] Compared with the prior art solutions, the advantages of the present invention are:

[0013] 1) Compared with the existing aerial image detection methods, the present invention proposes to use the knowledge distillation method to perform aerial image detection, which reduces the amount of network parameters and memory and improves the accuracy and robustness of student network detection.

[0014] 2) The present invention proposes an aerial image target detection method based on cross-level feature distillation. First, the query module performs loss on the corresponding target scale, which allows the student network to better learn the knowledge from the pre-trained teacher network. Secondly, cross-level feature distillation is performed on targets of different scales, and the multi-layer feature information of the pre-trained teacher network is used to guide the single-layer feature information of the student network, so that the semantic information between the pre-trained teacher network and the student network is more matched, which can effectively improve the effect of aerial image detection, while also reducing the amount of parameters and calculations, and reducing the network complexity. This method is novel and has certain application prospects in the deployment of satellite platforms. BRIEF DESCRIPTION OF THE DRAWINGS

[0015] Figure 1 It is a flow chart of the aerial image target detection method based on cross-level feature distillation proposed in the present invention.

[0016] Figure 2 This is the framework diagram of the teacher-student learning network based on cross-level feature distillation proposed in the present invention.

[0017] Figure 3 It is a specific distillation process based on cross-level feature distillation proposed in the present invention. DETAILED DESCRIPTION

[0018] The embodiments of the present invention are further described in detail below in conjunction with the accompanying drawings.

[0019] Combination Figure 1 and Figure 3 The present invention provides an aerial image target detection method based on cross-level feature distillation, comprising the following steps:

[0020] Step S1, obtain the aerial image detection dataset, randomly divide it into a training dataset and a test dataset, cut the images in the above two datasets into M×N blocks one by one, M and N represent the width and height of the image respectively, and the sizes of M and N are both between 800 and 1024. The training dataset is used to train the teacher network and the student network, and the test dataset is used to verify the student network after distillation, and then go to step S2.

[0021] Step S2: input the training data set into the teacher network for pre-training, obtain the pre-trained teacher network, and proceed to step S3.

[0022] Step S3: construct a distillation basic model based on the pre-trained teacher network and student network, and construct a teacher-student learning network framework based on cross-level feature distillation based on the distillation basic model (such as Figure 2 ):

[0023] The pre-trained teacher network includes a teacher backbone network, a teacher multi-scale fusion network and a teacher detection network. The teacher backbone network and the teacher multi-scale fusion network use HRNet32 and FPN respectively, and the teacher detection network includes a classification head, a regression head and a query head.

[0024] The student network includes a student backbone network, a student multi-scale fusion network and a student detection network. The student backbone network and the student multi-scale fusion network use HRNet18 and FPN respectively. The student detection network includes a classification head, a regression head and a query head.

[0025] Step S4, input the training data set into the teacher-student learning network framework, extract the features of each layer of the pre-trained teacher network and the teacher heat map of the query branch, and extract the features of each layer of the student network and the student heat map of the query branch, set a threshold to suppress the background pixels and low-confidence pixels in the teacher heat map and the student heat map, and obtain the pre-trained teacher network mask map and the student network mask map accordingly, and construct a query loss function of the student network mask map and the target truth mask map; construct a query module based on the pre-trained teacher network mask map and the student network mask map, perform cross-level feature distillation on targets of different scales based on the query module, and use the multi-layer feature information of the pre-trained teacher network to guide the single-layer feature information of the student network; fix the parameters in the pre-trained teacher network, and continuously update the training parameters of the student network by constructing the loss function of the student network, and finally obtain the trained student network. The specific steps are as follows:

[0026] The acquisition process of the pre-trained teacher network mask map and the student network mask map is as follows:

[0027] The targets are divided into four categories according to different scales: tiny targets, small targets, medium targets, and large targets. Tiny targets are defined as targets with a size of 10 to 20, small targets are defined as targets with a size of 20 to 50, medium targets are defined as targets with a size of 50 to 96, and large targets are defined as targets with a size greater than 96. The pre-trained teacher network generates four fusion feature layers with high to low resolutions {T 1 ,T 2 ,T 3 ,T 4}, and its corresponding teacher heat map is The existence probabilities of tiny targets, small targets, medium targets, and large targets are predicted respectively. To suppress background interference, a threshold is set on each teacher heat map to retain the target pixel value and suppress the background pixel value and the pixel value with low confidence. Teacher heat map The size is H l ×W l ,s l The final target pixel value is:

[0028]

[0029] Among them, (m,n) represents the coordinates of any pixel point in the teacher heat map.

[0030] The final pre-trained teacher network obtained by the pre-trained teacher network has a mask image of different target scales: Similarly, the student network mask images of different target scales obtained by the student network are

[0031] The query loss function of the student network mask map and the target truth mask map is constructed as follows:

[0032] The true value mask map of the target is calculated using the true value box. Assume that the true value box of target a on the k-th layer feature map is in, and represents the center position of the true value box of target a on the k-th layer feature map, represents the width of the truth box, Represents the height of the truth box, θ a Indicates the rotation angle of the ground truth box.

[0033] Calculate the center position of the truth box The distance D from any position (x, y) on the feature map k [x][y] is used to encode, and the calculation formula is:

[0034]

[0035] D k [x][y] is greater than d k Small is set to 1, otherwise it is set to 0, d k is the minimum scale of the preset anchor box of the k-th layer feature map.

[0036] Target ground truth mask for:

[0037]

[0038] Student Network Mask Map The corresponding target query loss function L query for:

[0039]

[0040] Among them, L FL It is the Focal Loss function, which is specially designed to deal with the foreground-background imbalance problem in dense target training. Represents the target truth mask map on the l-th layer feature map.

[0041] The specific process of constructing the query module based on the pre-trained teacher network mask graph and the student network mask graph is as follows:

[0042] The mask maps of different scales of the pre-trained teacher network and the student network are combined to obtain the target area. The mask map of the pre-trained teacher network of the lth scale and the mask map of the student network of the lth scale are taken respectively. The combination process of a single query module is:

[0043]

[0044] Among them, Q l Represents a single query module at each scale.

[0045] F assign (·,·) is used to combine the pre-trained teacher network mask graph and the student network mask graph. The combination process is:

[0046]

[0047] A single query module constitutes a query module group Q: Q = {Q 1 ,Q 2 ,Q 3 ,Q 4}, Q 1 Indicates the area where all tiny targets exist, Q 2 Indicates the area where all small targets exist, Q 3 represents all target areas, Q 4 Indicates the area where all large targets exist.

[0048] The specific operation of performing cross-level feature distillation on objects of different scales based on the query module is as follows:

[0049] The characteristics of each layer of the pre-trained teacher network are The characteristics of each layer of the student network are The dimensions of each layer of the student network are transformed, and the characteristics of each layer of the student network after transformation are defined as For the converted i-th layer student network features, this layer feature learning comes from the query-based module Q of the first i layers of the pre-trained teacher network i The feature of the i-th layer is the distillation loss L layer_i Defined as:

[0050]

[0051] i represents the number of student layers, and j represents the number of teacher layers; It is used to scale the size of each layer of the converted student network features to match the size of each layer of the pre-trained teacher network features; M(Q i ) represents the size conversion function, which is used to convert the query module Qi The size is converted into pre-trained teacher network features size; D(·,·) represents the distance between the features of each layer of the pre-trained teacher network and the features of each layer of the corresponding student network; the total loss function of distillation L cross_layer Defined as the sum of distillation losses at each layer:

[0052]

[0053] The total loss function construction process of the student network is as follows:

[0054] The parameters of the student network are continuously updated through the loss function. The cross-level feature distillation loss is added to the student network. The total loss function L of the student network is as follows:

[0055] L=L task +λL cross_layer

[0056] L task represents the original detection loss of the student network, including classification loss, regression loss, and target query loss. cross_layer represents the total distillation loss, λ is a balance coefficient used to balance the original detection loss and the total distillation loss, and then proceeds to step S5.

[0057] Step S5: Use the trained student network to detect all the block images on the test data set, merge the detection results on the block images, and finally output the categories and regression positions of all targets in the test data set in the aerial images, and finally obtain the accuracy of the trained student network.

Claims

1. A method for aerial image target detection based on cross-level feature distillation, characterized in that: The following steps are involved: Step S1, obtain an aerial image detection dataset, randomly divide it into a training dataset and a test dataset, cut the images in the above two datasets into M×N block images one by one, M and N represent the width and height of the image respectively, and the sizes of M and N are both between 800 and 1024. The training dataset is used to train the teacher network and the student network, and the test dataset is used to verify the student network after distillation, and then go to step S2; Step S2, input the training data set into the teacher network for pre-training, obtain the pre-trained teacher network, and proceed to step S3; Step S3: construct a distillation basic model based on the pre-trained teacher network and student network, and construct a teacher-student learning network framework based on cross-level feature distillation based on the distillation basic model, as follows: The pre-trained teacher network includes the teacher backbone network, the teacher multi-scale fusion network and the teacher detection network. The teacher backbone network and the teacher multi-scale fusion network use HRNet32 and FPN respectively. The teacher detection network includes a classification head, a regression head and a query head. The student network includes a student backbone network, a student multi-scale fusion network and a student detection network. The student backbone network and the student multi-scale fusion network use HRNet18 and FPN respectively. The student detection network includes a classification head, a regression head and a query head. Go to step S4; Step S4, input the training data set into the teacher-student learning network framework, extract the features of each layer of the pre-trained teacher network and the teacher heat map of the query branch, and extract the features of each layer of the student network and the student heat map of the query branch, set a threshold to suppress the background pixels and low-confidence pixels in the teacher heat map and the student heat map, and obtain the pre-trained teacher network mask map and the student network mask map accordingly, and construct a query loss function of the student network mask map and the target true value mask map; construct a query module based on the pre-trained teacher network mask map and the student network mask map, perform cross-level feature distillation on targets of different scales based on the query module, and use the multi-layer feature information of the pre-trained teacher network to guide the single-layer feature information of the student network; fix the parameters in the pre-trained teacher network, continuously update the training parameters of the student network by constructing the loss function of the student network, and finally obtain a trained student network, The acquisition of the pre-trained teacher network mask map and the student network mask map is as follows: The targets are divided into four categories according to different scales, and the pre-trained teacher network generates four teacher heat maps with high to low resolutions, which are: Predict the existence probabilities of micro-targets, small targets, medium targets, and large targets respectively; set a threshold s on each teacher's heat map l To suppress background pixels and pixels with low confidence, the teacher heat map The size is H l ×W l , l is the type label of the target scale, and its final target value pixel is: Where (m,n) represents the coordinates of any pixel in the teacher heat map; The final pre-trained teacher network obtained by the pre-trained teacher network has a mask image of different target scales: Similarly, the student network mask images of different target scales obtained by the student network are The query loss function of the student network mask map and the target truth mask map is constructed as follows: The true value mask map of the target is calculated using the true value box. Assume that the true value box of target a on the k-th layer feature map is in, represents the center position of the true value box of target a on the k-th layer feature map, represents the width of the truth box, Represents the height of the truth box, θ a Indicates the rotation angle of the truth box; Calculate the center position of the truth box The distance D from any position (x, y) on the feature map k [x][y] is used to encode, and the calculation formula is: D k [x][y] is greater than d k Small is set to 1, otherwise it is set to 0, d k is the minimum scale of the preset anchor box of the k-th layer feature map; Target ground truth mask for: Student Network Mask Map The corresponding target query loss function L query for: Among them, L FL It is the Focal Loss function, which is specially designed to deal with the foreground-background imbalance problem in dense target training. Represents the target truth mask map on the l-th layer feature map; Go to step S5; Step S5: Use the trained student network to detect all the block images on the test data set, merge the detection results on the block images, and finally output the categories and regression positions of all targets in the test data set in the aerial images, and finally obtain the accuracy of the trained student network.

2. The method for aerial image target detection based on cross-level feature distillation according to claim 1, characterized in that: The query module built based on the pre-trained teacher network mask graph and student network mask graph is as follows: The mask maps of different scales of the pre-trained teacher network and the student network are combined to obtain the target area. The mask map of the pre-trained teacher network of the lth scale and the mask map of the student network of the lth scale are taken respectively. The combination process of a single query module is: Among them, Q l Represents a single query module at each scale; F assign (·,·) is used to combine the pre-trained teacher network mask graph and the student network mask graph. The combination process is: A single query module constitutes a query module group Q: Q = {Q 1 ,Q 2 ,Q 3 ,Q 4 }, Q 1 Indicates the area where all tiny targets exist, Q 2 Indicates the area where all small targets exist, Q 3 represents all target areas, Q 4 Indicates the area where all large targets exist.

3. The method for aerial image target detection based on cross-level feature distillation according to claim 1, characterized in that: Based on the query module, cross-level feature distillation of objects of different scales is performed as follows: The characteristics of each layer of the pre-trained teacher network are The characteristics of each layer of the student network are The dimensions of each layer of the student network are transformed, and the characteristics of each layer of the student network after transformation are defined as For the converted i-th layer student network features, this layer feature learning comes from the query-based module Q of the first i layers of the pre-trained teacher network i The feature of the i-th layer is the distillation loss L layer_i Defined as: i represents the number of student layers, and j represents the number of teacher layers; It is used to scale the size of each layer of the converted student network features to match the size of each layer of the pre-trained teacher network features; M(Q i ) represents the size conversion function, which is used to convert the query module Q i The size is converted into pre-trained teacher network features size; D(·,·) represents the distance between the features of each layer of the pre-trained teacher network and the features of each layer of the corresponding student network; the total loss function of distillation L cross_layer Defined as the sum of distillation losses at each layer:

4. The method for aerial image target detection based on cross-level feature distillation according to claim 1, characterized in that: The total loss function of the student network is constructed as follows: The parameters of the student network are continuously updated through the loss function. The cross-level feature distillation loss is added to the student network. The total loss function L of the student network is as follows: L=L task +λL cross_layer L task represents the original detection loss of the student network, including classification loss, regression loss, and target query loss. cross_layer represents the total loss of distillation, and λ is the balance coefficient used to balance the original detection loss and the total loss of distillation.

Images