Image classification method based on full attention network structure search

Abstract

The present invention discloses an image classification algorithm based on full attention network structure search. First, a staged search space is designed, and different self-attention operations are selected in each stage of the network; then, a self-supervised search method is used to search, and the weight parameters and structure parameters inside the network are updated. When the self-supervised search stage is completed, the structure parameters are retained and used as the initial values ​​of the supervised search stage; finally, a supervised search method is used to search, and the weight parameters and structure parameters inside the network are updated, and the optimal full attention network is obtained according to the structure parameters. The search method proposed in the present invention can search for a high-performance full attention network structure while ensuring the required search efficiency. Experimental results on image classification tasks show that the searched full attention network model is superior to the most advanced architecture, while greatly reducing the number of parameters.

Description

Technical Field

[0001] The present invention relates to the field of image classification in computer vision, and more specifically, to a method for efficiently designing a full-attention network structure for image classification tasks based on a full-attention network structure search method. Background Art

[0002] In recent years, significant progress has been made in the design of attention network structures. Self-attention operations have become an important part of neural networks due to their ability to capture long-range feature dependencies and content-based parameter learning mechanisms. Neural networks composed entirely of self-attention operations have achieved good results in various computer vision tasks. Manually designing an excellent full-attention neural network structure is a challenging and complex task that requires a lot of prior knowledge and experience, as well as a lot of time and resources for experimental verification, which has greatly slowed down the development of attention networks.

[0003] Network structure search (NAS) provides a way to solve the above problems. NAS is the process of automating the design of neural network structures. Its goal is to automatically search for a network structure with better performance than the manually designed model from a given search space under certain prior knowledge. The NAS method can not only improve the performance of the model, but also free human experts from the tedious task of designing network structures. In recent years, many important advances have been made in the research of NAS. At present, the search methods of network structures are mainly divided into three types: reinforcement learning-based methods, evolutionary methods, and gradient-based differentiable network structure search methods. They mainly study how to use different search strategies to improve the efficiency of the search and the performance of the final network. Compared with the NAS method based on reinforcement learning and evolutionary methods, the gradient-based differentiable network structure search method converges faster and its optimization target is more flexible.

[0004] Existing NAS methods are not suitable for directly searching for full-attention networks. First, existing NAS methods usually adopt a cell-based search space. In the network structure obtained by the search, the cell structure of the shallow and deep layers of the network is the same. This is not suitable for self-attention network search because the effect of the self-attention operation is different at different stages of the network. In addition, existing NAS methods usually use classification tasks as supervision for structural search. Classification tasks require the model to pay more attention to local areas related to label information without considering the content correlation between distant pixels. However, the self-attention model focuses on capturing long-distance content dependencies between pixels to learn rich image representations. Therefore, it is not appropriate to directly use existing NAS methods to search for full-attention network structures. Summary of the invention

[0005] In order to solve the problems in the prior art, the present invention provides an image classification method based on full attention network structure search, which solves the problem of inaccurate full attention network search directly using the existing NAS method in the prior art.

[0006] The present invention is achieved through the following technical solutions:

[0007] An image classification method based on full attention network structure search, comprising the following steps:

[0008] Design a stage-wise search space where each stage of the network chooses a different self-attention operation;

[0009] Use a self-supervised search method to search, input the image into the network model in the phased search space, update the network internal weight parameters and structure parameters, and when the self-supervised search phase is completed, retain the structure parameters and use them as the initial values ​​of the supervised search phase;

[0010] Use the supervised search method to search, input the image into the network model in the staged search space, update the internal weight parameters and structural parameters of the network, and obtain the optimal full attention network based on the structural parameters.

[0011] In the network structure of the staged search space, the first layer is a fixed local self-attention operation, the last two layers are the average pooling layer and the classification layer, and the remaining intermediate part consists of five stages. In the second and fourth stages, there is a fixed pooling operation that halves the spatial size of the feature map and doubles the number of channels. Each stage has three searchable layers, and for each searchable layer, it is necessary to select the operation with the best performance from many candidate operations. The candidate operation consists of 7 self-attention operations, including a non-local self-attention operation and 6 local self-attention operations with different hyperparameters. The hyperparameters of the local self-attention operation include the spatial range and the number of heads, where the spatial range is selected from 3, 5 and 7, and the number of heads is selected from 4 and 8. In summary, the staged search space contains 15 searchable layers, each layer is selected from 7 candidate operations, and the search space contains 7 15 possible structure.

[0012] In the self-supervised search stage, a self-supervised search method based on the contextual autoregression task is designed; the contextual autoregression task refers to randomly masking multiple areas of the input image and training the network to predict the content information of the missing parts. An encoder-decoder structure is used to extract the features of the input image and reconstruct the missing image content. Then, this task is used to search for the full attention network. The full attention network is used as a feature encoder to extract features from the input image; the network contains two types of learnable parameters: the weight parameter w of the self-attention operation and the structural parameter a corresponding to each candidate operation; the image dataset is divided into two independent sets, represented by DatasetA and DatasetB respectively; the weight parameters are optimized using the DatasetA dataset, and the structural parameters are optimized using DatasetB; the L1 loss is used as the loss function, which is defined as follows:

[0013]

[0014] Where M is the number of pixels, p i is the input pixel, y i is the true value; then, the differentiable structure search method is used to alternately optimize the weight parameter w and the structure parameter a; in an iterative manner, the gradient descent is used on the DatasetA dataset To optimize the weight parameters, and then use gradient descent on the DatasetB dataset to optimize the structural parameters; when the self-supervised search phase is completed, store the structural parameters and use them as initial values ​​for the supervised search phase.

[0015] In the supervised search stage, the differentiable structure search method is used to search on the image classification dataset; the structure parameters obtained in the self-supervised search stage are used as initial values, and the gradient descent method is used to alternately optimize the structure parameters α and the weight parameters w; the cross entropy loss is used as the loss function, which is defined as:

[0016]

[0017] Where K is the number of categories, p k is the predicted probability of belonging to the kth class, y k is the category label; after the supervised search process is completed, the structural parameters are sorted, and the operation corresponding to the maximum value of the structural parameter is selected to obtain the final architecture.

[0018] The beneficial effect of the present invention is that the search method proposed by the present invention can search for a high-performance full-attention network structure while ensuring the required search efficiency. Experimental results on image classification tasks show that the searched full-attention network model outperforms the most advanced architecture while greatly reducing the number of parameters. BRIEF DESCRIPTION OF THE DRAWINGS

[0019] Figure 1 Network structure search flow chart;

[0020] Figure 2 Macroscopic network structure diagram of the search space;

[0021] Figure 3 The specific configuration of the candidate operation; where k represents the spatial range and h represents the number of heads in the local self-attention operation;

[0022] Figure 4 The best full attention network structure diagram obtained by search;

[0023] Figure 5 Performance comparison of the method of the present invention and existing advanced networks on CIFAR-10 and ImageNet; (a) Experimental results on the CIFAR dataset; (b) Experimental results on the ImageNet dataset. DETAILED DESCRIPTION

[0024] In order to make the technical solution of the present invention clearer, the present invention is further described below in conjunction with the accompanying drawings.

[0025] This paper proposes a full-attention network structure search method to efficiently design a full-attention network structure for image classification tasks. First, a staged search space is designed in which different self-attention operations can be selected at each stage of the network. Then, in order to efficiently find the optimal network structure from the search space, a new search strategy is proposed, which combines self-supervised search and supervised search to search for efficient full-attention network structures for image classification tasks. The method flowchart can be found in Figure 1 , the following is an introduction to the specific details of the search method:

[0026] 1. Phased Search Space

[0027] This paper proposes a phased search space for searching full attention networks. Figure 2 The macroscopic network structure of the search space is shown, and the number of layers and input dimensions of each stage of the network are shown in the figure. In the network structure of the search space, the first layer is a fixed local self-attention operation, and the last two layers are the average pooling layer and the classification layer. The remaining middle part consists of five stages. There is a fixed pooling operation in the second and fourth stages, which halves the spatial size of the feature map and doubles the number of channels. Each stage has three searchable layers. For each searchable layer, the operation with the best performance needs to be selected from many candidate operations. The candidate operations consist of 7 self-attention operations, including one non-local self-attention operation and 6 local self-attention operations with different hyperparameters. Their configurations are listed in Figure 3The hyperparameters of the local self-attention operation include the spatial range and the number of heads. The spatial range can be 3, 5, and 7, and the number of heads can be 4 and 8. In summary, our search space contains 15 searchable layers, each layer can choose from 7 candidate operations, and the search space contains 7 15 possible structure.

[0028] 2. Search strategy combining self-supervised search and supervised search

[0029] The first step is to prepare the dataset.

[0030] In the image classification task, the present invention selects the CIFAR-10 image classification dataset and the ImageNet image classification dataset to test and compare the performance of the method. The CIFAR-10 dataset contains images of 10 categories, each category has 6,000 images, a total of 60,000 images, which are divided into 50,000 training images and 10,000 test images. In the CIFAR dataset, the spatial resolution of each image is 32×32. The ImageNet 2012 dataset contains images of 1,000 categories, including 1.28 million training images and 50,000 verification images. For the ImageNet images, we crop them to 224×224 size.

[0031] In the second step, the search is performed using a self-supervised search method.

[0032] A training set was constructed by randomly selecting 100 classes from the 1000 original classes of ImageNet, and the images were resized to a resolution of 32×32. The training set was divided into two equal subsets, one for optimizing the network weight parameters and the other for optimizing the architecture parameters. During the self-supervised search process, the images were input into the network model in the staged search space, and the network parameters were optimized using the stochastic gradient descent (SGD) optimizer, where the weight parameter of the optimizer was set to 0.0003, the momentum decay rate was set to 0.9, and the initial learning rate of the network was 0.025. The structural parameters were optimized using the Adam optimizer with a learning rate of 3×10 -5 , the weight decays to 0.001. After 20 rounds of iterations, the search ends and the structural parameters are saved.

[0033] The third step is to use supervised search method to search.

[0034] The structural parameters obtained in the first stage are used as initialization values ​​to perform supervised search on the CIFAR-10 dataset. The training images of CIFAR-10 are randomly divided into two groups, each containing 25,000 images. One group is used to optimize the network weight parameters, and the other group is used to optimize the structural parameters. During the supervised search process, the images of the classification dataset are input into the network model of the staged search space, and the weight parameters are optimized using the SGD optimizer, with an initial learning rate of 0.025, a momentum of 0.9, and a weight decay of 0.0003. The weight parameters are optimized using the Adam optimizer, with a learning rate of 1×10 -4 , the weight decays to 0.001. After 50 rounds of iterations, the search ends and the optimal full attention network is obtained based on the structural parameters. The network obtained by the search is as follows Figure 4 shown.

[0035] The fourth step is to train and test the network structure on the CIFAR-10 and ImageNet datasets.

[0036] First, the network is trained on CIFAR-10 and the model is optimized using the SGD optimizer, where the weight parameter of the optimizer is set to 4×10 -4 , the momentum decay rate is set to 0.9, the initial learning rate is 0.04, and it gradually decays to 0 according to the cosine decay rule. After forward conduction and reverse conduction, the network parameter weights are updated once. After 500 rounds of iteration, the trained network can be obtained. During the test, the CIFAR-10 test set images are input into the network model to obtain the test results, such as Figure 5 (a) shown.

[0037] The network is then trained on ImageNet and the model is optimized using the SGD optimizer, where the weight parameter of the optimizer is set to 3×10 -5 , the momentum decay rate is set to 0.9, the initial learning rate is 0.04, and it gradually decays to 0 according to the cosine decay rule. After forward conduction and reverse conduction, the network parameter weights are updated once. After 300 rounds of iterations, the trained network can be obtained. During the test, the test set images of ImageNet are input into the network model to obtain the test results, such as Figure 5 (b) as shown.

[0038] The present invention compares the classification results of the searched model with those of the existing advanced models. From the experimental results, it can be seen that the method of the present invention greatly improves the accuracy of the searched full attention network structure in the classification task while ensuring the search efficiency.

Claims

1. An image classification method based on full attention network structure search, characterized in that: The following steps are involved: Design a staged search space in which each stage of the network selects a different self-attention operation; in the network structure of the staged search space, the first layer is a fixed local self-attention operation, the last two layers are an average pooling layer and a classification layer, and the remaining middle part consists of five stages, and there is a fixed pooling operation in the second and fourth stages, which halves the spatial size of the feature map and doubles the number of channels; each stage has three searchable layers, and for each searchable layer, it is necessary to select the operation with the best performance from multiple candidate operations; The candidate operations consist of 7 self-attention operations, including one non-local self-attention operation and 6 local self-attention operations with different hyperparameters. The hyperparameters of the local self-attention operation include the spatial range and the number of heads. The spatial range is 3, 5, and 7, and the number of heads is 4 and 8. In summary, the staged search space contains 15 searchable layers, each layer selects from 7 candidate operations, and the search space contains 7 15 possible structures; Use a self-supervised search method to search, input the image into the network model in the phased search space, update the network internal weight parameters and structure parameters, and when the self-supervised search phase is completed, retain the structure parameters and use them as the initial values ​​of the supervised search phase; Use the supervised search method to search, input the image into the network model in the staged search space, update the internal weight parameters and structural parameters of the network, and obtain the optimal full attention network based on the structural parameters.

2. The image classification method based on full attention network structure search according to claim 1 is characterized in that: In the self-supervised search stage, a self-supervised search algorithm based on the contextual autoregression task is designed; the contextual autoregression task refers to randomly masking multiple areas of the input image and training the network to predict the content information of the missing parts; an encoder-decoder structure is used to extract the features of the input image and reconstruct the missing image content; then, this task is used to search the full attention network; the full attention network is used as a feature encoder to extract features from the input image; the network contains two types of learnable parameters: the weight parameter w of the self-attention operation and the structural parameter a corresponding to each candidate operation; the image dataset is divided into two independent sets, represented by DatasetA and DatasetB respectively; the weight parameters are optimized using the DatasetA dataset, and the structural parameters are optimized using DatasetB; the L1 loss is used as the loss function, which is defined as follows: Where M is the number of pixels, p i is the input pixel, y i is the true value; then, the differentiable structure search method is used to alternately optimize the weight parameter w and the structure parameter a; in an iterative manner, the gradient descent is used on the DatasetA dataset To optimize the weight parameters, and then use gradient descent on the DatasetB dataset to optimize the structural parameters; when the self-supervised search phase is completed, store the structural parameters and use them as initial values ​​for the supervised search phase.

3. The image classification method based on full attention network structure search according to claim 1 is characterized in that: In the supervised search stage, the differentiable structure search method is used to search on the image classification dataset; the structure parameters obtained in the self-supervised search stage are used as initial values, and the gradient descent method is used to alternately optimize the structure parameters α and the weight parameters w; the cross entropy loss is used as the loss function, which is defined as: Where K is the number of categories, p k is the predicted probability of belonging to the kth class, y k is the category label; after the supervised search process is completed, the structural parameters are sorted, and the operation corresponding to the maximum value of the structural parameter is selected to obtain the final architecture.

Images