RESOLVING INTRA-CLASS IMBALANCE FOR GAN-BASED IMAGE AUGMENTATION

Abstract

Advanced machine learning and deep learning techniques have increasingly improved accuracy of image classification. Most existing studies have investigated the data imbalance problem among classes to further enhance classification accuracy. However, less attention has been paid to data imbalance within every single class. In this work, we present AC-GAN (Actor-Critic Generative Adversarial Network), a data augmentation framework that explicitly considers heterogeneity of intra-class data. AC-GAN exploits a novel loss function to weigh the impacts of different subclasses of data in a class on GAN training. It hence can effectively generate fake data of both majority and minority subclasses, which help train a more accurate classifier. We use defect detection as an example application to evaluate our design. The results demonstrate that the intra-class distribution of fake data generated by our AC-GAN can be more similar to that of raw data. With balanced training for various subclasses, AC-GAN enhances classification accuracy for no matter uniformly or non-uniformly distributed intra-class data.