CN118429733A - Multi-head attention-driven kitchen garbage multi-label classification method - Google Patents

Abstract

The invention discloses a multi-head attention-driven kitchen garbage multi-label classification method, which comprises the following steps: constructing a kitchen waste multi-label classification data set, wherein the kitchen waste multi-label classification data set comprises a plurality of images of different categories of kitchen waste, and the image labels comprise one or more categories; constructing a multi-head attention-driven graph rolling light-weight network model, wherein the model comprises a feature extraction module, a multi-head attention module and a dynamic graph rolling module; the feature extraction module extracts features from an input image, sends the features into the multi-head attention module to strengthen a category sensing area of a feature map, and sends the features into the dynamic map convolution module to self-adaptively capture the category sensing area; training the graph roll lightweight network model by using the constructed kitchen waste multi-label classification data set; and finally, performing multi-label classification on the kitchen waste image to be predicted by using the classification model obtained through training. The invention can reduce the performance loss caused by the reduction of model parameters, enhance the recognition capability and improve the multi-label classification effect.

Description

A multi-label classification method for kitchen waste driven by multi-head attention

Technical Field

The present invention relates to the field of image processing, and in particular to a multi-head attention driven kitchen waste multi-label classification method.

Background technique

In recent years, many advanced classification algorithms have been proposed as a means to improve the efficiency of garbage classification with deep learning. However, most of the currently available garbage datasets are designed based on the recognition of domestic garbage, and there is a lack of relevant research on the classification of kitchen waste in real scenes. In addition, kitchen waste images often contain multiple categories, which is a typical multi-label image classification in the field of computer vision. As one of the tasks in the field of computer vision, multi-label image classification is mainly aimed at accurately predicting all the categories contained in the image. Since in the real world, multiple objects usually appear at the same time, this is more in line with human cognitive common sense.

From the perspective of multi-label classification of kitchen waste, real scenes have high requirements for the real-time performance of the model. Due to the complex image background, diverse object categories, and correlation between object labels, the current methods for dealing with this problem often require more complex network models if they want to obtain higher classification accuracy, and the basic deep convolutional network cannot guarantee classification accuracy. Therefore, it is of great practical value and significance to study intelligent kitchen waste multi-label efficient classification algorithms in real scenes.

Summary of the invention

The present invention provides a multi-head attention driven kitchen waste multi-label classification method, which reduces the performance loss caused by the reduction of model parameters while further enhancing the recognition ability and improving the multi-label classification effect.

In order to achieve the above technical objectives, the present invention adopts the following technical solutions:

A multi-head attention driven kitchen waste multi-label classification method, including:

Construct a multi-label classification dataset for kitchen waste, including several images of different categories of kitchen waste, and the label of each image includes one or more categories;

Construct a multi-head attention-driven graph convolution lightweight network model, including a lightweight feature extraction module, a multi-head attention module and a dynamic graph convolution module; the feature extraction module extracts feature maps from the image of the input model, and the extracted feature maps are sent to the multi-head attention module for processing to strengthen the category perception area of the feature maps. The feature maps with enhanced category perception areas are then sent to the dynamic graph convolution module for processing to adaptively capture the category perception area and output the predicted category;

Use the constructed kitchen waste multi-label classification dataset to train the graph convolution lightweight network model;

Finally, the trained model, namely the kitchen waste multi-label classification model, is used to perform multi-label classification on the kitchen waste images to be predicted.

Furthermore, based on the different category distributions of kitchen waste in different seasons, the kitchen waste multi-label classification dataset is obtained by selecting kitchen waste images of multiple different category combinations from different time periods of a year.

Furthermore, the lightweight feature extraction module uses a lightweight backbone network ShufflenetV2.

Furthermore, the multi-head attention module includes a first fully connected layer, a scaled dot product attention submodule, a second fully connected layer, a Dropout layer and a normalization layer;

The first fully connected layer takes the input feature map Dimensionality reduction and conversion to feature map ;in, Represent the length, width, and number of channels of the image respectively;

The scaled dot product attention submodule adopts a multi-head attention mechanism, and each head takes the feature map As key and value, query uses a set of learnable parameters, calculated as:

In the formula, represents the output of the scaled dot-product attention submodule using the multi-head attention mechanism, represents feature concatenation, is the additional weight matrix, The first The output of the head, They are the query, key, value and corresponding weight matrix of the scaled dot product attention submodule input Multiply them together to get; is the scaling factor; Respectively The weight matrix to be learned is the weight matrix The weight in the i-th dimension;

The second fully connected layer, the Dropout layer, and the normalization layer further process the output of the scaled dot product attention submodule, which is expressed as:

In the formula, represents point-by-point addition, , , Respectively represent the processing of the second fully connected layer, Dropout layer and normalization layer; Feature map representing the output of the multi-head attention module.

Furthermore, the dynamic graph convolution module includes a static graph convolution layer and a dynamic graph convolution layer;

The dynamic graph convolution layer outputs the static graph convolution layer Processing, expressed as:

In the formula, is the correlation matrix of the dynamic graph convolutional layer, and is the state update weight of the dynamic graph convolution layer, it's for Activation function, for Activation function; is to transform the feature map Its global representation The global representation is obtained by concatenation. The output of the static graph convolution layer Pooling, 1*1 one-dimensional convolution and activation function are performed to obtain it; C represents the total number of categories for multi-label classification of kitchen waste, and D ₁ represents the dimension of the output feature map H of the static graph convolution layer.

Furthermore, the static graph convolution layer processes the input feature graph, which is expressed as:

In the formula, is the feature map output by the multi-head attention module, is the output feature map of the static graph convolution layer, which consists of features corresponding to C categories, namely ; is the activation function, is the correlation matrix of the static graph convolutional layer, and is the state update weight of the static graph convolution layer.

The present invention discloses a multi-head attention driven kitchen waste multi-label classification method, which uses a lightweight network to optimize the graph convolution classification model, and introduces a multi-head attention mechanism to reduce the loss of feature information, capture feature information at different levels, enhance the feature extraction capability of the backbone network in complex scenarios, reduce the performance loss caused by the reduction of model parameters, and further use a dynamic graph convolution module to achieve adaptive capture of semantic perception areas, further enhance recognition capabilities, and improve multi-label classification effects. Compared with existing kitchen waste classification technologies, it has the following advantages:

(1) To overcome the shortcomings of the traditional GCN method, which uses ordinary deep convolutional networks as the feature extraction backbone network with low accuracy and Transformer as the feature extraction backbone network with large number of model parameters, the present invention adopts ShuffleNetV2 as the feature extraction backbone network and optimizes the network model to achieve lightweight.

(2) The present invention designs a multi-head attention module and a dynamic graph convolution module to optimize the lightweight graph convolution classification network, thereby reducing the performance loss caused by the reduction of model parameters, further enhancing the recognition ability and improving the multi-label classification effect.

(3) It has strong practicality and generalization ability. The present invention not only achieves good results in the current multi-label classification of kitchen waste, but also obtains excellent multi-label classification accuracy on the MS-COCO and VOC 2007 datasets.

BRIEF DESCRIPTION OF THE DRAWINGS

FIG1 is a flow chart of a multi-label classification method for kitchen waste driven by multi-attention according to an embodiment of the present invention.

FIG2 is a schematic diagram of a multi-head attention module according to an embodiment of the present invention.

FIG. 3 shows the relationship between different backbone network parameters and accuracy according to an embodiment of the present invention.

FIG4 shows the influence of the fusion of different modules on the classification mAP according to an embodiment of the present invention.

FIG5 shows the influence of the fusion of different modules on the classification of APs of various categories according to an embodiment of the present invention.

Detailed ways

The following is a detailed description of an embodiment of the present invention. This embodiment is based on the technical solution of the present invention, and provides a detailed implementation method and a specific operation process to further explain the technical solution of the present invention.

The technical solution of the present invention is to realize multi-label classification of kitchen waste images based on a multi-head attention-driven graph convolution lightweight network. The Python programming language can be used for experiments, and the C/C++ programming language can also be used for engineering applications.

The present invention provides a multi-head attention driven graph convolution lightweight network multi-label classification method. The processing flow is shown in Figure 1, which includes the following steps:

Step 1: Build a multi-label classification dataset MLKW for kitchen waste collected in real scenarios, including the collection and annotation of the dataset;

The data described in the embodiment is collected on the assembly line of the kitchen waste sorting center, and the 15,994 kitchen waste images obtained complete the production of the MLKW data set. MLKW is closer to the real scene and has a strong engineering application significance. In order to reflect the characteristics of kitchen waste in different time periods, considering the different distribution of kitchen waste categories in different seasons, the present invention selects images from different time periods of the year, including 3771 images in spring, 4735 images in summer, 2130 images in autumn and 4358 images in winter. By cleaning and screening the image data, a multi-label kitchen waste data set containing eight categories, a total of 3107 fully annotated images, and the rest are annotated is finally obtained, and the data set is divided into a training set and a test set in a ratio of 8:2. At the same time, verification was performed on the PASCAL VOC 2007 dataset and MS-COCO dataset, which are widely used in the field of multi-label classification. The VOC 2007 dataset contains a total of 9,963 images and covers 20 common categories. The MS-COCO dataset contains a training set (82,081 images) and a validation set (40,504 images), a total of 122,581 images, covering 80 common categories, and each image has approximately 2.9 category labels.

Step 2: Build a multi-head attention-driven graph convolution lightweight network model, including a lightweight feature extraction module, a multi-head attention module, and a dynamic graph convolution module.

1. Feature extraction module.

In this embodiment, the lightweight backbone network ShufflenetV2 is used as the feature extraction network of the GCN module to achieve lightweight, i.e., GCLN. However, in order to adapt to the lightweight model of the present invention, the original resolution of the image, 3,256×2,724, is adjusted to 448×448. Before the model is trained, ShuffleNetV2_x1_0, which has been pre-trained on the ImageNet1K dataset, is first used as the backbone network.

2. Multi-head attention module.

The feature map extracted by the feature extraction module is sent to the multi-head attention module to strengthen the category perception area of the feature map.

The purpose of the multi-head attention module (MHA) is to capture the category area of interest, through the features extracted by the backbone network , W, H, and D represent the length, width, and number of channels of the image, respectively. After dimensionality reduction, it is transformed into , which is used as the key and value of the multi-head attention mechanism. For the query, a set of learnable parameters is used to complete the global review of the image features, see Figure 2. In MHA, scaled dot-product attention is used as the calculation core, and its calculation formula is:

(1)

Among them: Q, K, V are the input q', k', v' through the fully connected layer and the corresponding weight matrix Multiply them together to get . is the scaling factor, The dimension of is the same as the dimension of the initial input q, k, v. See Figure 2.

Different from the multi-head attention mechanism in Vision Transformer, this paper introduces residual connection to reduce information loss by redesigning it. The calculation process is as follows: map the q, k, and v matrices to different subspace transformations to generate multiple heads, and perform a scaled dot product attention mechanism on each head to capture information in different subspaces and dimensions. Among them:

(2)

(3)

Among them: W is the additional weight matrix, and Concat represents feature concatenation.

After multi-head attention, considering the overfitting problem in the training process, Dropout layer and normalization are added to alleviate it. In addition, residual connection is used to reduce information loss. The calculation formula is as follows:

(4)

Where: V _c represents the category of interest, c represents the number of categories, It means point-by-point addition, Q is q' and Multiply to.

3. Dynamic convolution module.

The feature map output by the multi-head attention module is fed into the dynamic graph convolution module for final classification.

This example uses a multi-head attention mechanism to obtain the category of interest In order to make the graph convolution have a specific label dependency for each image, a dynamic graph convolution module (Dynamic GCN) is designed.

By passing Vc to static GCN and dynamic GCN in sequence. A single-layer static GCN is simply defined as ,in , activation function for ,in, The correlation matrix, is the state update weight. Then use dynamic GCN. The correlation matrix It is a feature matrix affected by image features. This method can effectively alleviate the overfitting problem. The output of dynamic GCN , the specific calculation process is:

(5)

in, for Activation function, for Activation function, is the state update weight, To construct a weight for the dynamic correlation matrix Ad, is obtained by taking H and its global representation This is obtained by concatenating the global average pooling and a conv layer in sequence. Formally, defined as:

(6)

Step 3: Use the constructed kitchen waste multi-label classification dataset to train the graph convolution lightweight network model.

During the training process, the dynamic graph convolution module (DGCN) uses the nonlinear activation function LeakyReLU, and sets the slope to 0.2. Some hyperparameters in the training process are as follows: the optimizer uses SGD, its momentum is 0.9, and the weight decay coefficient is 0.0001. The batch size is 16, the initial learning rate of the multi-head attention module (MHA) and the dynamic graph convolution module (DGCN) is 0.5, and the initial learning rate of the backbone network is 0.05. A total of 50 epochs are trained, and the learning rate is updated with a decay coefficient of 0.1 at the 30th and 40th epochs respectively. Ubuntu 16.04LTS system or above is required, and the system environment requires Pytorch 1.6, Python3.6 or above. The hardware platform includes graphics card Tesla V100, CUDA11.0, CuDNN8.0.2GPUs, and the CPU memory is not less than 64G, and the solid-state drive is not less than 512G.

Step 4: Use the trained model, i.e., the kitchen waste multi-label classification model, to perform multi-label classification on the kitchen waste image to be predicted.

To facilitate understanding of the technical effects of the present invention, the following comparison is provided between the present invention and the traditional method:

Table 1 and Figure 3 show the comparison of different lightweight backbone networks in different methods. The lightweight backbone networks include ShuffleNetV2, MobileNetV3, and EfficientNet, and the multi-label image classification algorithms include ML-GCN, ADD-GCN, and MHA-GCN. The main comparison is the average precision (mAP) and parameter amount (Params) of different backbone networks between different algorithms. As can be seen from Table 1, under the same backbone, the present invention is superior to other methods in all indicators. As can be seen from Figure 2, by slightly increasing the number of parameters of the model, a large performance improvement is obtained. The performance of the present invention on the MLKW dataset is at least 8.6% higher than that of ML-GCN and at least 4.8% higher than that of ADD-GCN, which shows the effectiveness of the method in this paper.

Table 2 shows the experimental results of the present invention and other methods on the VOC2007 dataset. It can be seen that the AP values on multiple categories are superior to other methods, and compared with ResNet101, the mAP is increased to 94.0%. It is worth noting that the present invention is designed for the kitchen waste dataset and has the same performance as ML-GCN on VOC2007, which also proves the generalization and scalability of the method.

Table 3 is the experimental results of other methods of the present invention on the MS-COCO dataset. Table 3 shows the performance of MHA-GCN on MS-COCO, as well as the comparison results with ResNet101, SRTN, ML-GCN and other methods. Obviously, the method MHA-GCN of the present invention outperforms other methods in multiple indicators such as mAP, OF1, CF1, etc. In addition, compared with ResNet101, the overall performance of MHA-GCN is improved by 6.1%, which further demonstrates the advantages of the method of this paper.

In order to further verify the effectiveness of the multi-head attention module MHA and the dynamic graph convolution module GCN of the method of the present invention, the present invention verifies the influence of the fusion of different modules on the classification mAP and the influence of the fusion of different modules on the AP of each category of classification on the MLKW data, see Figure 4 and Figure 5 respectively.

The above embodiments are preferred embodiments of the present application. Ordinary technicians in this field can also make various changes or improvements on this basis. Without departing from the overall concept of the present application, these changes or improvements should fall within the scope of protection required by the present application.

Claims (6)

1. A multi-head attention-driven kitchen garbage multi-label classification method is characterized by comprising the following steps of:

Constructing a kitchen waste multi-label classification data set, wherein the multi-label classification data set comprises a plurality of images of kitchen waste in different categories, and the labels of each image comprise one or more categories;

constructing a multi-head attention-driven graph rolling light-weight network model, wherein the model comprises a light-weight characteristic extraction module, a multi-head attention module and a dynamic graph rolling module; the feature extraction module extracts feature images of an input model, the extracted feature images are sent to the multi-head attention module to be processed so as to strengthen category sensing areas of the feature images, and the feature images after strengthening the category sensing areas are sent to the dynamic image convolution module to be processed so as to adaptively capture the category sensing areas and output prediction categories;

training the graph roll lightweight network model by using the constructed kitchen waste multi-label classification data set;

Finally, a model obtained through training, namely a kitchen waste multi-label classification model, is used for multi-label classification of the kitchen waste image to be predicted.

2. The multi-head attention-driven kitchen waste multi-label classification method according to claim 1, wherein the kitchen waste multi-label classification data set is obtained by selecting kitchen waste images with a plurality of different types and combinations from different time periods of one year based on different types and distributions of kitchen waste in different seasons.

3. The multi-head attention driven kitchen waste multi-tag classification method of claim 1, wherein the lightweight feature extraction module uses a lightweight backbone network ShufflenetV.

4. The multi-head attention driven kitchen waste multi-label classification method according to claim 1, wherein the multi-head attention module comprises a first full-connection layer, a zoom dot product attention sub-module, a second full-connection layer, a Dropout layer and a normalization layer;

The first full connection layer inputs the characteristic diagram Dimension reduction conversion into feature map; Wherein,Representing the length, width, and number of channels of the image, respectively;

The zoom dot product attention submodule adopts a multi-head attention mechanism, and each head is used for mapping the characteristic diagram As the key and the value, the query adopts a set of parameters which can be learned, and the calculation formula is as follows:

；

；

；

In the method, in the process of the invention, Representing the output of a scaled dot product attention sub-module employing a multi-headed attention mechanism,The feature stitching is represented and is performed,In order to attach the weight matrix to the vehicle,Attention module for dot product scalingThe output of the individual head is provided with,Query, key, value and corresponding weight matrix, respectively, of the scaled dot product attention submodule inputMultiplying to obtain; is a scaling factor; Respectively the first The weight matrix to be learned of each head is the weight matrixWeights in the ith dimension;

The second fully-connected layer, dropout layer, and normalization layer further process the output of the scaled dot product attention sub-module, expressed as:

；

In the method, in the process of the invention, Representing a point-by-point addition,、、Respectively representing the treatment of a second full connection layer, a Dropout layer and a normalization layer; and the characteristic diagram represents the output of the multi-head attention module.

5. The multi-head attention driven kitchen waste multi-label classification method according to claim 1, wherein the dynamic graph convolution module comprises a static graph convolution layer and a dynamic graph convolution layer;

Output of the dynamic graph convolution layer to the static graph convolution layer The process is expressed as:

；

In the method, in the process of the invention, Is the correlation matrix of the dynamic graph convolution layer, andIs the state update weight of the dynamic graph convolutional layer,Is thatThe function is activated and the function is activated,Is thatActivating a function; By combining characteristic diagrams And global representation thereofObtained by concatenation, where global representsOutput through static graph convolutional layersCarrying out pooling and 1*1 one-dimensional convolution and activating functions to obtain; c represents the total number of categories of multi-label classification of kitchen waste, and D ₁ represents the dimension of an output characteristic diagram H of a static diagram convolution layer.

6. The multi-head attention driven kitchen waste multi-label classification method according to claim 5, wherein the static map convolution layer processes an input feature map, expressed as:

；

In the method, in the process of the invention, Is a characteristic diagram output by the multi-head attention module,The output characteristic diagram of the static diagram convolution layer consists of C types of corresponding characteristics, namely；In order to activate the function,Is the correlation matrix of the convolution layer of the static diagramIs the state update weight of the static graph convolutional layer.