A two-stage long tail learning method

Abstract

The application provides a two-stage long tail learning method, comprising: performing two-stage multi-round iterative training on an image classification model by using a training set of image classification with a long tail distribution to obtain a trained image classification model, wherein: first stage learning comprises: determining a first sharpness perception gradient by using a head class set and a loss function, and determining a second sharpness perception gradient by using a tail class set and the loss function, and updating parameters of the image classification model according to the first sharpness perception gradient and the second sharpness perception gradient; second stage learning comprises: determining a first original gradient by using the head class set and the loss function, and determining a second sharpness perception gradient by using the tail class set and the loss function, and updating the parameters of the image classification model according to the first original gradient and the second sharpness perception gradient, and the method improves the generalization ability of the model as a whole.

Description

Technical Field

[0001] This invention relates to the field of neural network technology, specifically to the field of deep learning under long-tail distribution, and more specifically, to a two-stage long-tail learning method. Background Technology

[0002] Long-tail distribution [1][2] A data distribution characterized by a small number of classes (generally called head classes) occupying a large portion of the sample size, while a large number of classes (generally called tail classes) occupy a small portion, exhibiting a significant imbalance. Deep learning models applied to long-tail distributions are called long-tail learning models.

[0003] Data collected from the real world often exhibits a long-tail distribution. For example, in wildlife image classification and recognition, the number of images containing rare animals is often far less than the number of images containing common animals such as tigers and elephants. Due to its practical significance, long-tail learning has received increasing attention in recent years.

[0004] Traditional long-tail learning methods are often based on resampling or category-sensitive learning. The former modifies the dataset, generating balanced data through undersampling or oversampling; the latter modifies the loss function, changing the weights and margins of different classes within the loss function. However, traditional methods often exhibit overfitting on the tail classes, severely limiting the model's generalization ability.

[0005] Smoothness regularization can effectively solve the overfitting problem. [3][4] Sharpness-Perceived Minimization (SAM) [5] It is a typical smoothness regularization method that minimizes the sharpness of the loss function while minimizing the loss value, so that the model converges to a flat local or global minimum point, thereby effectively improving the generalization ability of the model and providing a promising direction for solving the overfitting problem.

[0006] However, directly applying the sharpness-perceived minimization method to long-tail learning is not ideal, because the head class has a larger sample size, and the influence of the head class in its perturbation term will dominate, affecting the generalization ability on the tail class.

[0007] The information for the above references is as follows:

[0008] [1]Wanli Ouyang,Xiaogang Wang,Cong Zhang,and Xiaokang Yang.Factors infinetuning deep model for object detection with long-tail distribution.InCVPR,pages 864–873,2016.

[0009] [2]Songyang Zhang,Zeming Li,Shipeng Yan,Xuming He,and JianSun.Distribution alignment:A unified framework for long-tail visualrecognition.In CVPR,pages 2361–2370,2021.

[0010] [3]Bingyi Kang,Saining Xie,Marcus Rohrbach,Zhicheng Yan,Albert Gordo,Jiashi Feng,and Yannis Kalantidis.Decoupling representation and classifierfor long-tailed recognition.In ICLR,2020.

[0011] [4]Shaden Alshammari,Yu-Xiong Wang,Deva Ramanan,and Shu Kong.Long-tailed recognition via weight balancing.In CVPR,pages 6897–6907,2022.

[0012] [5]Pierre Foret,Ariel Kleiner,Hossein Mobahi,and BehnamNeyshabur.Sharpness-aware minimization for efficiently improvinggeneralization.In ICLR,2021.

[0013] It should be noted that the background information presented here is only for illustrating relevant information about the present invention to aid in understanding the technical solution of the present invention, and does not imply that the relevant information is necessarily prior art. The relevant information was submitted and disclosed together with the present invention, and should not be considered prior art unless there is evidence that the relevant information was disclosed before the filing date of the present invention. Summary of the Invention

[0014] Therefore, the purpose of this invention is to overcome the shortcomings of the prior art and provide a two-stage long-tail learning method.

[0015] The objective of this invention is achieved through the following technical solution:

[0016] According to a first aspect of the present invention, a two-stage long-tail learning method is provided, comprising: performing two-stage multi-round iterative training on an image classification model using a training set of image classification with a long-tail distribution to obtain a trained image classification model, wherein each round of training includes: obtaining a batch of samples from the training set, classifying the batch of samples into a head class set or a tail class set according to whether the labels contained in the batch of samples belong to a preset head class or a tail class set; determining whether the model training has met preset requirements, if not, performing the first stage of learning in the current round, if yes, performing the second stage of learning in the current round, wherein: the first stage of learning includes: determining a first sharpness-perceived gradient using the head class set and a loss function, and determining a second sharpness-perceived gradient using the tail class set and a loss function, and updating the parameters of the image classification model according to the first sharpness-perceived gradient and the second sharpness-perceived gradient; performing the second stage of learning includes: determining a first original gradient using the head class set and a loss function, and determining a second sharpness-perceived gradient using the tail class set and a loss function, and updating the parameters of the image classification model according to the first original gradient and the second sharpness-perceived gradient.

[0017] Optionally, the samples include sample images and labels, where the labels indicate the category to which the corresponding sample image belongs. The first stage of learning includes: inputting the sample images from the batch into an image classification model to obtain a first predicted value; determining a first original gradient corresponding to the head class set based on the first predicted value, label, and loss function of the sample images in the head class set; calculating a perturbation term for the head class using the first original gradient based on a sharpness-aware minimization technique; adding the parameters of the image classification model to the perturbation term of the head class to obtain a first perturbation parameter; using the first perturbation parameter to predict the sample images of the head class set to obtain a second predicted value; and calculating the gradient based on the second predicted value, label, and loss function to obtain... The image classification model is calculated as follows: First sharpness-aware gradient for the head class set; Second original gradient for the tail class set based on first predicted values, labels, and loss functions for sample images in the tail class set; Perturbation term for the tail class is calculated using the second original gradient based on sharpness-aware minimization technique; Second perturbation parameter is obtained by adding the parameters of the image classification model to the perturbation term for the tail class; Third predicted value is obtained by predicting sample images of the tail class set using the second perturbation parameter; Gradient is calculated based on the third predicted value, label, and loss function to obtain the second sharpness-aware gradient for the tail class set; Finally, the parameters of the image classification model are updated based on the total gradient calculated from the first and second sharpness-aware gradients.

[0018] Optionally, the second stage of learning includes: inputting sample images from the batch into the image classification model to obtain a first predicted value; determining the first original gradient corresponding to the head class set based on the first predicted value, label, and loss function corresponding to the sample images in the head class set; determining the second original gradient corresponding to the tail class set based on the first predicted value, label, and loss function corresponding to the sample images in the tail class set; calculating the perturbation term of the tail class using the second original gradient based on the sharpness-aware minimization technique; adding the parameters of the image classification model to the perturbation term of the tail class to obtain the second perturbation parameter; using the second perturbation parameter to predict the sample images of the tail class set to obtain a third predicted value; calculating the gradient based on the third predicted value, label, and loss function to obtain the second sharpness-aware gradient corresponding to the tail class set; and updating the parameters of the image classification model based on the total gradient calculated from the first original gradient and the second sharpness-aware gradient.

[0019] Optionally, the perturbation term for the head class is calculated as follows:

[0020]

[0021] Where, ∈ head (w t ) represents the perturbation term for the head class, which is a vector that includes the perturbation component corresponding to each parameter component in the image classification model; ρ head The neighborhood size of the head class is a hyperparameter, ρ.tail >ρ head >0; This represents the first original gradient calculated based on the first predicted value, label, and loss function of the sample image using the head class set. represents the original gradient of the entire batch calculated based on the first predicted value, label, and loss function of the sample images in that batch, and ‖·‖2 represents the Euclidean norm.

[0022] Optionally, the tail-class perturbation term is calculated as follows:

[0023]

[0024] Where, ∈ tail (w t ) represents the perturbation term for the tail class, which is a vector containing the perturbation component corresponding to each parameter component in the image classification model; ρ tail The neighborhood size of the tail class is a hyperparameter, ρ. head >0; This represents the second primitive gradient calculated based on the first predicted value, label, and loss function of the sample image from the tail class set. represents the original gradient of the entire batch calculated based on the first predicted value, label, and loss function of the sample images in that batch, and ‖·‖2 represents the Euclidean norm.

[0025] Optionally, update the parameters of the image classification model as follows:

[0026] w t+1 =w t -η(g head +g tail )

[0027] Among them, w t This represents the parameters of the image classification model used in the current round, where η represents the learning rate, and g... head In the first learning phase, g represents the first sharpness perception gradient. head In the second learning phase, g represents the first original gradient. tail This represents the second sharpness perception gradient.

[0028] According to a second aspect of the present invention, an image classification method is provided, comprising: obtaining a trained image classification model obtained according to the method of the first aspect; obtaining an image to be classified, and classifying the image to be classified using the trained image classification model.

[0029] According to a third aspect of the present invention, a computer program product is provided, comprising a computer program / instructions that, when executed by a processor, implement the steps of the methods described in the first and / or second aspects.

[0030] According to a fourth aspect of the present invention, an electronic device is provided, comprising: one or more processors; and a memory for storing executable instructions; wherein the one or more processors are configured to implement the steps of the methods of the first aspect and / or the second aspect by executing the executable instructions. Attached Figure Description

[0031] The embodiments of the present invention will be further described below with reference to the accompanying drawings, wherein:

[0032] Figure 1 This is a flowchart illustrating a two-stage long-tail learning method according to an embodiment of the present invention. Detailed Implementation

[0033] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be understood that the specific embodiments described herein are merely illustrative and are not intended to limit the invention.

[0034] As mentioned in the background section, directly applying the sharpness-perceived minimization method to long-tail learning is undesirable because the head class has a larger sample size, and its influence dominates the perturbation term, affecting the generalization ability on the tail classes. To address this, the inventors considered designing a novel long-tail learning method that, while solving the tail class overfitting problem through sharpness-perceived minimization, further optimizes the model's generalization ability on the tail classes without excessively limiting the model's performance on the head classes, thereby improving the overall model performance. Therefore, the inventors designed a two-stage long-tail learning method. In the first stage of learning, both the sharpness-perceived gradient of the head classes (i.e., the first sharpness-perceived gradient) and the sharpness-perceived gradient of the tail classes (i.e., the second sharpness-perceived gradient) are used to optimize the parameters of the classification model. Although the influence of the head class will dominate, a second stage of learning will be performed after the first stage of learning is completed. At this time, only sharpness-aware minimization is considered for the tail class, while the head class does not use sharpness-aware minimization but considers the regular gradient (i.e. the first original gradient), thereby further improving the generalization ability of the tail class.

[0035] According to one embodiment of the present invention, a two-stage long-tail learning method is provided, see [link to relevant documentation]. Figure 1 The process includes steps S1 to S5. To better understand this invention, each step will be described in detail below with reference to specific embodiments.

[0036] Step S1: Obtain the image classification model and initialize the parameters of the image classification model.

[0037] According to one embodiment of the present invention, the image classification model can employ existing classification models, such as ResNet, DenseNet, or VGGNet. Alternatively, a user-defined model can be used, for example, an image classification model constructed using convolutional layers, residual layers, attention layers, linear layers, and classification layers, or combinations thereof. This embodiment of the present invention does not impose any limitations on this. After obtaining the model, the parameters of the image classification model can be randomly initialized and used as the parameters of the model for the first round of training, denoted as w0.

[0038] Step S2: Extract the training set from the image classification dataset with long-tail distribution properties, and determine the head class and tail class according to the number of samples in each category of the dataset.

[0039] According to one embodiment of the present invention, the image classification dataset with long-tail distribution properties can be an existing dataset, such as CIFAR-100-LT, CIFAR-10-LT, ImageNet-LT, or the iNaturalist dataset. The dataset includes multiple samples, each including a sample image and a label, where the label indicates the true category to which the corresponding sample image belongs. Of course, implementers can also construct their own datasets with long-tail distribution properties relevant to the desired application scenario, as long as the data distribution property of a small number of categories (generally called head classes) occupying a large number of samples and a large number of categories (generally called tail classes) occupying a small number of samples is satisfied. This embodiment of the present invention does not impose any limitations on this.

[0040] Step S3: Perform multiple rounds of the first stage of learning until the model training meets the preset requirements, and then proceed to the second stage of learning. The first stage of learning includes: obtaining a batch of samples extracted from the training set and a head class set or a tail class set obtained based on the batch of samples; determining the first sharpness-aware gradient using the head class set and the loss function; determining the second sharpness-aware gradient using the tail class set and the loss function; and updating the parameters of the classification model based on the first and second sharpness-aware gradients.

[0041] The preset requirements can be set according to the implementer's needs. According to one embodiment of the invention, this can be achieved based on the current training epoch (current epoch) and a preset training threshold. If the current epoch is less than the preset training threshold, the preset requirements are considered not met, and the first stage of learning is executed; if the current epoch is greater than or equal to the preset training threshold, the preset requirements are considered met, and the second stage of learning is executed. Alternatively, other methods can be used, such as observing whether the classification model has been trained to a preset convergence condition. If the convergence condition is met, the second stage of learning is executed for the current epoch and thereafter.

[0042] The head class set or tail class set can be defined based on the labels. According to one embodiment of the invention, samples are generally randomly selected from the dataset according to a preset ratio (e.g., 7:3 or 8:2) to form a training set and a test set. Alternatively, the implementer can divide it into a training set, a validation set, and a test set. Then, in each training round, a batch of samples is drawn from the training set. Each batch includes multiple samples, each containing a sample image and a label. If the label of a sample in this batch belongs to the head class, the sample is assigned to the head class set; otherwise, it is assigned to the tail class set. For example, a threshold η is set. thres For a category y in the dataset, if the number of samples corresponding to it is greater than η thres If a sample is classified as a head sample, it is classified as a tail sample; otherwise, it is classified as a tail sample. During the subsequent batch selection process, the class of each sample is checked to determine whether it belongs to the head or tail class, thus completing the classification.

[0043] According to one embodiment of the present invention, the first stage of learning includes:

[0044] Input the sample images from this batch into the image classification model w t The first predicted value was obtained;

[0045] Based on the first predicted value, label, and loss function corresponding to the sample images in the head class set, determine the first original gradient corresponding to the head class set.

[0046] Based on the sharpness-aware minimization technique, the perturbation term of the head class is calculated using the first original gradient, and is expressed as:

[0047]

[0048] Where, ∈ head (w t ) represents the perturbation term for the head class, which is a vector that includes the perturbation component corresponding to each parameter component in the image classification model; ρ head The neighborhood size of the head class is a hyperparameter, ρ. head >0; This represents the first original gradient calculated based on the first predicted value, label, and base loss function of the sample image using the head class set. This represents the original gradient of the entire batch calculated based on the first predicted value, label, and basic loss function of the sample images in this batch, and ||·|2 represents the Euclidean norm.

[0049] The parameters of the image classification model are added to the perturbation term of the head class to obtain the first perturbation parameter w′. t =w t +∈ head (wt ), using the first perturbation parameter w′ t A second predicted value is obtained by predicting the sample images of the head class set. The gradient is then calculated based on the second predicted value, the label, and the base loss function to obtain the first sharpness-perceived gradient corresponding to the head class set.

[0050] Based on the first predicted value, label, and base loss function corresponding to the sample images in the tail class set, determine the second original gradient corresponding to the tail class set.

[0051] Based on the sharpness-aware minimization technique, the perturbation term of the tail class is calculated using the second primitive gradient, and is expressed as:

[0052]

[0053] Where, ∈ tail (w t ) represents the perturbation term for the tail class, which is a vector containing the perturbation component corresponding to each parameter component in the image classification model; ρ tail The neighborhood size of the tail class is a hyperparameter, ρ. head >0; This represents the second primitive gradient calculated based on the first predicted value, label, and base loss function of the sample image from the tail class set. This represents the original gradient of the entire batch calculated based on the first predicted value, label, and basic loss function of the sample images in this batch, and ||·|2 represents the Euclidean norm.

[0054] The parameters of the image classification model are added to the perturbation term of the tail class to obtain the second perturbation parameter w″. t =w t +∈ tail (w t ), using the second perturbation parameter w″ t A third predicted value is obtained by predicting the sample images of the tail class set. The gradient is then calculated based on the third predicted value, the label, and the base loss function to obtain the second sharpness-aware gradient corresponding to the tail class set.

[0055] The parameters of the image classification model are updated based on the total gradient calculated from the first and second sharpness-perceived gradients: w t+1 =w t -η(g head +g tail ), where η represents the learning rate.

[0056] The above ρ head ρ tail The recommended value range is 0.05 to 0.30, and it is recommended that ρ... tail>ρ head The value of η can range from 0.001 to 0.2.

[0057] In the first stage of learning described above, the main goal is to help head classes escape saddle points. The loss function can be expressed as follows:

[0058]

[0059] in, Let w represent the total loss function for the first stage. t This represents the parameters used by the image classification model at round t (the parameters before being updated at round t). This represents the loss function used to calculate the first sharpness-sensing gradient. This represents the loss function used to calculate the second sharpness-perceived gradient. As those skilled in the art will know, during training, the goal of updating the model's parameters is to minimize the value of the total loss function. In this embodiment of the invention, the goal of updating the parameters of the image classification model in the first stage is to minimize... The value of .

[0060] It should be noted that the above and They rely on the same basic loss function k, only differing in the parameters and / or the object used to calculate the loss. Specifically, It uses the sample images of the entire batch as input parameters w to the model. t The predicted values ​​are used to calculate the loss for the entire batch. It uses sample images from the head class set as inputs to the model's parameters w. t The predicted values ​​are used to calculate the loss of the head class set. It uses sample images from the tail class set as inputs to the model's parameters w. t The predicted values ​​are used to calculate the loss of the tail class set. The first perturbation parameter w′ is input using the sample images of the head class set. t =w t +∈ head (w t The predicted values ​​obtained are used to calculate the loss of the head class set. The second perturbation parameter w" is input using the sample images of the tail class set. t =w t +∈ tail (w t The predicted values ​​obtained are used to calculate the loss of the tail class set.

[0061] Step S4: Perform multiple rounds of second-stage learning, including: obtaining a batch of samples extracted from the training set and a head class set or tail class set obtained based on the batch of samples; determining the first original gradient using the head class set and the loss function; determining the second sharpness-aware gradient using the tail class set and the loss function; and updating the parameters of the image classification model based on the first original gradient and the second sharpness-aware gradient.

[0062] According to one embodiment of the present invention, the second stage of learning includes:

[0063] Input the sample images from this batch into the image classification model to obtain the first predicted value;

[0064] The first original gradient corresponding to the head class set is determined based on the first predicted value, label, and loss function of the sample images in the head class set.

[0065] Based on the first predicted value, label, and loss function corresponding to the sample images in the tail class set, determine the second original gradient corresponding to the tail class set;

[0066] Based on the sharpness-aware minimization technique, the perturbation term of the tail class is calculated using the second primitive gradient;

[0067] The parameters of the image classification model are added to the perturbation term of the tail class to obtain the second perturbation parameter. The second perturbation parameter is used to predict the sample images of the tail class set to obtain the third prediction value. The gradient is calculated based on the third prediction value, the label and the loss function to obtain the second sharpness perception gradient corresponding to the tail class set.

[0068] The parameters of the image classification model are updated based on the total gradient calculated from the first original gradient and the second sharpness-perceived gradient.

[0069] The difference between the second stage and the first stage is that the second stage updates the parameters of the image classification model based on the total gradient calculated from the first original gradient and the second sharpness-aware gradient. It is not necessary to calculate the perturbation term, the first perturbation parameter, and the first sharpness-aware gradient for the head class. The implementation details of the rest are the same as those of the first stage, and will not be repeated here.

[0070] In the second stage of learning described above, the main goal is to assist head classes in escaping saddle points. The loss function can be expressed as follows:

[0071]

[0072] in, This represents the total loss function for the second stage. The loss function used to calculate the first original gradient. In this embodiment of the invention, the goal of updating the parameters of the image classification model in the second stage is to minimize... The value of .

[0073] Step S5: After the learning of the second phase of the last round is completed, a trained image classification model is obtained.

[0074] According to an embodiment of the present invention, a practitioner may set the total number of training rounds, and after the number of trained rounds meets the requirement of the total number of training rounds, the trained image classification model is obtained. Alternatively, a final convergence determination condition may also be set, and when the learning process is in the second phase and the model has been trained to convergence, the trained image classification model is obtained.

[0075] According to an example of the present invention, a two-stage long-tail learning method is provided, comprising the steps of:

[0076] A1. Acquire a training set for image classification a loss function l, a batch size n, a learning rate η, a total number of training rounds T, and a neighborhood size ρ of head classes head , and a neighborhood size ρ of tail classes tail , and a phase transition factor γ, where ρ tail >ρ head >0, 0<γ<1, x i represents the i-th sample image, y i represents the label corresponding to the i-th sample image, and N represents the total number of samples in the training set;

[0077] A2. Acquire an image classification model, initialize the parameter w0 of the image classification model, and set the initial round t=0;

[0078] A3. When t<T, execute the following steps:

[0079] A3.1. If t≥γT holds, set the neighborhood size ρ of the head classes head to 0, enter the second phase of learning, and proceed to A3.2; otherwise, directly proceed to A3.2;

[0080] A3.2. Randomly sample one batch from the training set

[0081] A3.3. Divide into a head class batch and a tail class batch

[0082] A3.4. Calculate the gradient of the head class loss and the gradient of the tail class loss where w t represents the parameter of the model in the t-th round, represents the loss function value on the head class sample , represents the tail class sample The loss function value on;

[0083] A3.5 Calculate the perturbation term for the head class:

[0084]

[0085] in, represents the loss function value over the entire batch, and ||·|2 represents the Euclidean norm;

[0086] A3.6 Calculate the perturbation term for the tail class:

[0087]

[0088] A3.7 Calculate the gradient of the sharpness-perceived minimization loss function for the head class:

[0089]

[0090] Where, ∈ head (w t ) represents the perturbation term of the head class;

[0091] A3.8 Calculate the gradient of the sharpness-perceived minimization loss function for the tail class:

[0092]

[0093] Where, ∈ tail (w t ) represents the perturbation term of the tail class;

[0094] A3.9, Update Weights:

[0095] w t+1 =w t -η(g head +g tail )

[0096] A3.10, Increase t by 1:

[0097] A4. Returns the parameters w of the final trained model. t .

[0098] It should be noted that, due to ρ tail Set to 0 in the second stage. ∈ head (w t ) = 0, g head The gradient is actually equal to the head class loss. That is, the first original gradient; and thus, by setting ρ tail The value is used to switch from the first stage to the second stage.

[0099] According to one embodiment of the present invention, an image classification method is also provided, comprising: obtaining a trained image classification model obtained by a two-stage long-tail learning method according to the foregoing embodiments; obtaining an image to be classified; and classifying the image to be classified using the trained image classification model.

[0100] To verify the effectiveness of the method of the present invention, the inventors also conducted comparative experiments.

[0101] The datasets selected for the experiment were CIFAR-10-LT and CIFAR-100-LT.

[0102] The basic loss function forms (i.e., the loss function l mentioned earlier): CE, LDAM, LA, and VS, where:

[0103] CE stands for Cross-Entropy, which is the cross-entropy loss function.

[0104] LDAM stands for label-distribution-aware margin loss, which is a label distribution-aware margin loss function.

[0105] LA stands for logit adjustment, which is a loss function for adjusting the model output logistic value. It is a loss function that calculates the cross-entropy after adjusting the model output logistic value based on the class distribution.

[0106] VS stands for vector-scaling loss, which is the vector scaling loss function.

[0107] Experimental results

[0108]

[0109]

[0110] Among them, some methods do not include SAM or SSE-SAM, indicating that the model parameters are updated to minimize the loss value by directly calculating the original gradient based on the base loss function (which is an existing method).

[0111] Some lines of methods include SAM, which stands for sharpness-aware minimization (the existing method).

[0112] The method in some rows carries SSE-SAM, indicating the two-stage long-tail learning method of the present invention, in which the loss function form before SSE-SAM is used.

[0113] The experimental results show that, compared with existing methods, the model trained by the method of this invention has the highest accuracy in most cases, indicating that the method of this invention improves the generalization ability of the model as a whole.

[0114] In general, traditional long-tail learning methods often face overfitting issues on the tail classes, while smoothness regularization can effectively address this problem. This invention proposes a novel two-stage long-tail learning method based on sharpness-aware minimization, improving upon existing methods. Specifically, theoretical analysis shows that compared to stochastic gradient descent, sharpness-aware minimization enhances the model's ability to escape saddle points, thereby improving its generalization ability. Building on this, this invention focuses on saddle point escape for the head and tail classes respectively in the first and second training stages, demonstrating complementary advantages and thus improving the overall generalization ability of the model. Experiments show that the proposed method effectively improves the model's generalization ability.

[0115] It should be noted that although the steps are described in a specific order above, it does not mean that the steps must be executed in the above specific order. In fact, some of these steps can be executed concurrently, or even in a different order, as long as the required function can be achieved.

[0116] This invention can be a system, method, and / or computer program product. A computer program product may include a computer-readable storage medium having computer-readable program instructions loaded thereon for causing a processor to implement various aspects of the invention.

[0117] Computer-readable storage media can be tangible devices that hold and store instructions for use by an instruction execution device. Computer-readable storage media can be, for example, including but not limited to, electrical storage devices, magnetic storage devices, optical storage devices, electromagnetic storage devices, semiconductor storage devices, or any suitable combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), static random access memory (SRAM), portable compact disc read-only memory (CD-ROM), digital multifunction disc (DVD), memory sticks, floppy disks, mechanical encoding devices, such as punch cards or recessed protrusions storing instructions thereon, and any suitable combination thereof.

[0118] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or technical improvements to the embodiments in the market, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. A two-stage long-tail learning method, characterized in that, include: The image classification model is trained in two stages using a training set with a long-tailed distribution. Each training stage includes: A batch of samples is obtained from the training set. Based on whether the labels contained in the batch of samples belong to the preset head class or tail class, the batch of samples is assigned to the head class set or the tail class set. Determine whether the model training has met the preset requirements. If not, proceed to the first stage of learning in the current round; if yes, proceed to the second stage of learning in the current round. The first phase of learning includes: determining the first sharpness-perceived gradient using the head class set and loss function, and determining the second sharpness-perceived gradient using the tail class set and loss function, and updating the parameters of the image classification model based on the first and second sharpness-perceived gradients. The second phase of learning is performed, which includes: determining the first original gradient using the head class set and the loss function, and determining the second sharpness-aware gradient using the tail class set and the loss function, and updating the parameters of the image classification model based on the first original gradient and the second sharpness-aware gradient, wherein the sample images in the batch are input into the image classification model to obtain the first predicted value. The first original gradient corresponding to the head class set is determined based on the first predicted value, label, and loss function of the sample images in the head class set. Based on the first predicted value, label, and loss function corresponding to the sample images in the tail class set, determine the second original gradient corresponding to the tail class set; Based on the sharpness-aware minimization technique, the perturbation term of the tail class is calculated using the second primitive gradient; The parameters of the image classification model are added to the perturbation term of the tail class to obtain the second perturbation parameter. The second perturbation parameter is used to predict the sample images of the tail class set to obtain the third prediction value. The gradient is calculated based on the third prediction value, the label and the loss function to obtain the second sharpness perception gradient corresponding to the tail class set. The parameters of the image classification model are updated based on the total gradient calculated from the first original gradient and the second sharpness-perceived gradient.

2. The method according to claim 1, characterized in that, The samples include sample images and labels, where the labels indicate the category to which the corresponding sample image belongs. The first phase of learning includes: Input the sample images from this batch into the image classification model to obtain the first predicted value; The first original gradient corresponding to the head class set is determined based on the first predicted value, label, and loss function of the sample images in the head class set. Based on the sharpness-perceived minimization technique, the perturbation term of the head class is calculated using the first original gradient; The parameters of the image classification model are added to the perturbation term of the head class to obtain the first perturbation parameter. The first perturbation parameter is used to predict the sample images of the head class set to obtain the second predicted value. The gradient is calculated based on the second predicted value, the label and the loss function to obtain the first sharpness perception gradient corresponding to the head class set. Based on the first predicted value, label, and loss function corresponding to the sample images in the tail class set, determine the second original gradient corresponding to the tail class set; Based on the sharpness-aware minimization technique, the perturbation term of the tail class is calculated using the second primitive gradient; The parameters of the image classification model are added to the perturbation term of the tail class to obtain the second perturbation parameter. The second perturbation parameter is used to predict the sample images of the tail class set to obtain the third prediction value. The gradient is calculated based on the third prediction value, the label and the loss function to obtain the second sharpness perception gradient corresponding to the tail class set. The parameters of the image classification model are updated based on the total gradient calculated from the first and second sharpness-perceived gradients.

3. The method according to claim 2, characterized in that, The head-type perturbation term is calculated as follows: in, The perturbation term representing the head class is a vector that includes the perturbation component corresponding to each parameter component in the image classification model. The size of the neighborhood of the header class is a hyperparameter. >0; This represents the first original gradient calculated based on the first predicted value, label, and loss function of the sample image using the head class set. This represents the initial gradient of the entire batch, calculated based on the first predicted value, label, and loss function of the sample images in that batch. This indicates the search for the Euclidean norm.

4. The method according to claim 3, characterized in that, The tail-type perturbation term is calculated as follows: in, The perturbation term representing the tail class is a vector that includes the perturbation component corresponding to each parameter component in the image classification model. The size of the tail class's neighborhood is a hyperparameter. >0; This represents the second primitive gradient calculated based on the first predicted value, label, and loss function of the sample image from the tail class set. This represents the initial gradient of the entire batch, calculated based on the first predicted value, label, and loss function of the sample images in that batch. This indicates the search for the Euclidean norm.

5. The method according to claim 4, characterized in that, The parameters of the image classification model are updated in the first learning phase as follows: in, This represents the parameters of the image classification model used in the current round. Indicates the learning rate. This represents the first sharpness perception gradient. This represents the second sharpness perception gradient.

6. The method according to claim 4, characterized in that, In the second learning phase, the parameters of the image classification model are updated as follows: in, This represents the parameters of the image classification model used in the current round. Indicates the learning rate. This represents the first original gradient. This represents the second sharpness perception gradient.

7. An image classification method, comprising: Obtain a trained image classification model according to the method of any one of claims 1-6; Obtain the image to be classified, and use the trained image classification model to classify the image.

8. A computer program product comprising a computer program / instructions that, when executed by a processor, implement the steps of the method according to any one of claims 1-7.

9. A computer-readable storage medium, characterized in that, It stores a computer program that can be executed by a processor to implement the steps of the method according to any one of claims 1-7.

10. An electronic device, characterized in that, include: One or more processors; as well as Memory, wherein the memory is used to store executable instructions; The one or more processors are configured to implement the steps of the method according to any one of claims 1-7 by executing the executable instructions.

Images