BAGAN: Data Augmentation with Balancing GAN

Giovanni Mariani, Florian Scheidegger, Roxana Istrate, Costas Bekas, and Cristiano MalossiIBM Research – Zurich, Switzerland

Abstract

Image classification datasets are often imbalanced, char-acteristic that negatively affects the accuracy of deep-learning classifiers. In this work we propose balancingGAN (BAGAN) as an augmentation tool to restore bal-ance in imbalanced datasets. This is challenging becausethe few minority-class images may not be enough totrain a GAN. We overcome this issue by including duringthe adversarial training all available images of majorityand minority classes. The generative model learns usefulfeatures from majority classes and uses these to generateimages for minority classes. We apply class conditioningin the latent space to drive the generation process towardsa target class. The generator in the GAN is initializedwith the encoder module of an autoencoder that enablesus to learn an accurate class-conditioning in the latentspace. We compare the proposed methodology withstate-of-the-art GANs and demonstrate that BAGANgenerates images of superior quality when trained withan imbalanced dataset.

1 Introduction

The accuracy of image classification techniques can sig-nificantly deteriorate when the training dataset is im-balanced, i.e. when available data is not uniformlydistributed between the different classes. Imbalanceddatasets are common and a traditional approach to miti-gate this problem is to augment the dataset by introduc-ing additional minority-class images derived by applyingsimple geometric transformations to original images, e.g.rotations or mirroring. This augmentation approachmay disrupt orientation-related features when these arerelevant. In this work we propose a balancing genera-tive adversarial network (BAGAN) as an augmentationtool to restore the dataset balance by generating newminority-class images. Since these images are scarce inthe initial dataset, it is challenging to train a GAN forgenerating new ones. To overcome this issue, the pro-posed methodology includes in the adversarial trainingall data from minority and majority classes at once. This

IBM, the IBM logo, and ibm.com are trademarks or registeredtrademarks of International Business Machines Corporation inthe United States, other countries, or both. Other product andservice names might be trademarks of IBM or other companies.Submitted. Copyright 2018 by the author(s).

enables BAGAN to learn underlying features of the spe-cific classification problem starting from all images andthen to apply these features for the generation of newminority-class images. For example, let us consider theclassification of road traffic signs [1]. All warning signsshare the same external triangular shape. Once BAGANlearns to draw this shape from one of the signs, we canapply it for drawing any other one. Since BAGAN learnsfeatures starting from all classes whereas the goal is togenerate images for the minority classes, a mechanismto drive the generation process toward a desired classis needed. To this end, in this work we apply class con-ditioning on the latent space [2, 3]. We initialize thediscriminator and generator in the GAN with an autoen-coder. Then, we leverage this autoencoder to learn classconditioning in the latent space, i.e. to learn how theinput of the generative model should look like for differ-ent classes. Additionally, this initialization enables us tostart the adversarial training from a more stable pointand helps mitigating convergence problems arising withtraditional GANs [4, 5, 6, 7]. The main contributions ofthis work are:

• An overall approach to train GANs with an imbal-anced dataset while specifically aiming to generateminority-class images.

• An autoencoder-based initialization strategy thatenables us to a) start training the GAN from a goodinitial solution, and b) learn how to encode differentclasses in the latent space of the generator.

• An empirical evaluation of the proposed BAGANmethodology against the state of the art.

Experimental results empirically demonstrate that theproposed BAGAN methodology outperforms state-of-the-art GAN approaches in terms of variety and qualityof the generated images when the training dataset isimbalanced. In turn, this leads to a higher accuracy offinal classifiers trained on the augmented dataset.

2 Background

In recent years generative adversarial neural networks(GANs) [8, 9, 7] have been proposed as a tool to artifi-cially generate realistic images. The underlying idea is

1

arX

iv:1

803.

0965

5v2

[cs

.CV

] 5

Jun

201

8

to train a generative network in adversarial mode againsta discriminator network.

A well known problem of generative adversarial modelsis that while they learn to fool the discriminator they mayend up drawing one or few foolish examples. This prob-lem is known as mode collapse [8, 4, 5, 6]. In this work,our aim is to augment an imbalanced image classificationdataset to restore its balance. It is of paramount impor-tance that the augmented dataset is variable enough anddoes not include a continuously repeating example, thuswe need to avoid mode collapse. To this end, differentapproaches have been proposed. Possible solutions are:explicitly promoting image diversity in the generatorloss [10, 4], letting the generator predict future changesof the discriminator and adapt against these [11], letthe discriminator distinguish the different classes [2, 12],applying specific regularization techniques [5, 6], andcoupling GANs with autoencoders [4, 13, 14, 15].

In this work we apply the latter approach and cou-ple GAN and autoencoding techniques. The cited ap-proaches include additional modules in the GAN to em-bed an autoencoder all along the training. In the pro-posed BAGAN methodology we apply a more pragmaticapproach and use an autoencoder to initialize the GANmodules close to a good solution and far from mode col-lapse. Since our goal is to generate images specifically forthe minority classes, we train a generator that is control-lable in terms of the image class it draws, similarly to thestate-of-the-art ACGAN methodology [2]. Nonetheless,ACGAN is not specifically meant for imbalanced datasetsand turns to be flawed when targeting the generation ofminority-class images.

3 Motivating Example

State-of-the-art GANs are not suitable to deal with im-balanced datasets [16] and, to the best of our knowledge,the proposed BAGAN methodology is the first one tospecifically address this topic. Before going through de-tails of the proposed approach, let us demonstrate witha simple example why it is difficult to apply existingGAN techniques to tackle the problem at hand. Let usconsider the classification of handwritten digits, startingfrom an imbalanced version of the MNIST dataset [17]where we remove 97.5% of the available zeros from thetraining set.

A trivial idea would be to use a traditional GAN[8, 9, 18], train it by using all the available data, generatemany random samples, find the 0 instances and use thesefor augmenting the dataset. This approach cannot beapplied in general: if the generator G in the GAN istrained to fool the discriminator D by generating realisticimages, it will better focus on the generation of majorityclasses to optimize its loss function while collapsing awaythe modes related to the minority class. On the other

hand, training a GAN by using only the minority-classimages is not really an option because minority-classimages are scarce. In this example, after removing 97.5%of the zeros, we are left with about 150 minority-classimages. In general, it is difficult to train a GAN startingfrom a very little dataset, the GAN must have manyexamples to learn from [19].

A different approach is to train the GAN with ma-jority and minority classes jointly and to let the GANdistinguish between different classes explicitly. Duringtraining, the generator is explicitly asked to draw imagesof every class and to let the discriminator believe that thegenerated images are real images of the desired classes.While doing so, the generator is explicitly rewarded fordrawing realistic images of each class including the mi-nority classes. To the best of our knowledge, the onlymethod implementing this approach presented in liter-ature so far is ACGAN [2] where the generator inputcan be conditioned to draw a target class. In ACGAN,the discriminator has two outputs, one to discriminatebetween real and fake images X, and the other to classifyX in terms of its class c, Figure 1(a). During training,the generator is explicitly asked to draw images Xc foreach class c. Generator parameters are adjusted to max-imize the superposition of two components. The firstcomponent is the log-likelihood of generating an imageXc considered real by the discriminator. The secondcomponent is the log-likelihood of generating an imageXc that the discriminator associates with the class c. Weobserved that, when a dataset is imbalanced, these twocomponents become contradictory for the minority class.This can be explained as follows. Let us assume that ata point in time the generator converged into a solutionwhere it generates minority-class images with real-word

𝑋c1

cn

…

fakereal

Discriminator

(a) ACGAN discriminator.

fakec1

cn

…𝑋Discriminator

(b) BAGAN discriminator.

Figure 1: Discriminator architectures for ACGAN andBAGAN.

(a) ACGAN.

(b) Proposed BAGAN.

Figure 2: Ten zero-digit images generated with ACGANand the proposed BAGAN when trained with an imbal-anced version of MNIST where 97.5% of the zeros weredropped.

2

quality. These images would be indistinguishable bythe discriminator from the ones in the training dataset.Since in the training dataset minority-class images arescarce, when a minority-class image is passed to the dis-criminator during training it is most likely a fake image.To optimize its loss function the discriminator has toassociate the fake label to all minority-class images. Atthis point, the two generator objectives are contradictoryand the generator can either draw an image that looksreal or one that is representative of the minority class butcannot achieve these two goals at once. In turn, the gen-erator can be rewarded for drawing images that look realand are not representative of the target minority class.This fact deteriorates the quality of generated images.Images for the imbalanced MNIST example generated byACGAN for the 0 digit are shown in Figure 2(a). In thiswork we propose BAGAN that applies class conditioningas ACGAN but differs on the following points.

First, the BAGAN discriminator has a single outputthat returns either a problem-specific class label c or thelabel fake, Figure 1(b). The discriminator D is trainedfor associating to the images generated by G the labelfake, and to real images Xc their class label c. Thegenerator is trained to avoid the fake label and matchthe desired class labels. Since this is now defined as asingle objective rather than as a superposition of twoobjectives, by construction it cannot contradict itselfand the generator is never rewarded for the generationof images Xc that look real if the discriminator does notmatch them with the desired class label c.

Second, BAGAN couples GAN and autoencoding tech-niques to provide a precise selection of the class condi-tioning and to better avoid mode collapse. Images forthe imbalanced MNIST example generated by BAGANare of superior quality, Figure 2(b).

4 BAGAN

The proposed BAGAN methodology aims to generaterealistic minority-class images for an imbalanced dataset.It exploits all available information of the specific clas-sification problem by including in the BAGAN trainingmajority and minority classes jointly. GAN and autoen-

coding techniques are coupled to leverage the strengths ofthe two approaches. GANs generate high-quality imageswhereas autoencoders converge towards good solutioneasily [7]. Several authors suggest to couple GANs andautoencoders [4, 13]. Nonetheless these works are notdirectly meant to drive the GAN generative process to-wards specific classes. It is not easy to generalize them toenable the GAN to distinguish between different classes.As explained in the motivating example, in this work weapply class conditioning as suggested by Odena et al. [2]to embed class knowledge in BAGAN.

We apply a pragmatic use of autoencoders to initializethe GAN close to a good solution and far from modecollapse. Additionally, we apply the encoder part ofthe autoencoder to infer the distribution of the differentclasses in the latent space. The autoencoder-based GANinitialization is achievable by using the same networktopology in the autoencoder and GAN modules, Figures3(a) and 3(b). The decoding stage ∆ of the autoencodermatches the topology of the generator G. The encod-ing stage E of the autoencoder matches the topologyof the first layers of the discriminator De. In BAGAN,the knowledge in the autoencoder is transferred into theGAN modules by initializing the parameter weights ac-cordingly, Figure 3(b). To complete the discriminator, afinal dense layer Dd with a softmax activation functiontranslates the latent features into the probability thatthe image is fake or that it belongs to one of the problemclasses c1 — cn. When the GAN modules are initial-ized, a class-conditional latent vector generator is set upby learning the probability distribution of the imagesin the latent space for the different classes. Then, allthe weights in the generator and discriminator are finetuned by carrying out a traditional adversarial training,Figure 3(c). Overall, the BAGAN training approach isorganized in the three steps show in Figure 3: a) autoen-coder training, b) GAN initialization, and c) adversarialtraining.

Autoencoder training. The autoencoder is trainedby using all the images in the training dataset. Theautoencoder has no explicit class knowledge, it processesall images from majority and minority classes uncondi-tionally. In this work we apply l2 loss minimization forthe autoencoder training.

Autoencoder

E D

Encoder DecoderLatentvectorZOriginal

dataReconstructed

datal2 loss optimization

(a) Autoencoder training.

GAN

G = D

Generator

Latentvector

Z

Fakeimage

X

De=E Dc

Discriminatorfakec1

cn

…ImageX

Originaldata

E

c1instances

𝒩1

cninstances

E𝒩n

… Class-conditionallatent vectorgenerator

(b) GAN initialization.

GANGenerator

Latentvectors

Zc

Fakeimages

DDiscriminator

fakec1

cn

…

Class-conditionallatent vectorgenerator

G

Realimages

Originaldata Sparse categorical cross-entropy optimization

Classlabels

c

(c) GAN training.

Figure 3: The three training steps of the proposed BAGAN methodology.

3

GAN initialization. Differently from the autoen-coder, the generator G and the discriminator D haveexplicit class knowledge. During the adversarial training,G is asked to generate images for different classes, andD is asked to label the images either as fake or with aproblem-specific class label c. At the moment the GANis initialized, the autoencoder knowledge is transferredinto the GAN modules by initializing G with the weightsin the decoder ∆, and the first layers of the discriminatorDe with the weights of the encoder E, Figure 3(b). Thelast layer of the discriminator Dd is a dense layer witha softmax activation function and generates the finaldiscriminator output. The weights of this last layer areinitialized at random and learnt during the adversarialtraining.

The discriminator initialization is used simply to in-clude inD meaningful features that can help in classifyingimages. The initialization of the generator has a deeperreason. When adversarial training starts, the generatorG is equivalent to the decoder ∆. Thus a latent vectorZ input to the generator G is equivalent to a point inthe latent space of the autoencoder, i.e. Z can be seenas the output of E or the input of ∆. Thus, the encoderE maps real images into the latent space in use by G.We leverage this fact to learn a good class conditioningbefore to start the adversarial training, i.e. we definehow a latent vector Zc should look like for an image ofclass c.

We model a class in the latent space with a multivariatenormal distribution Nc = N (µc,Σc) with mean vector µc

and covariance matrix Σc. For each class c, we computeµc and Σc to match the distribution of Zc = E(Xc)considering all real images Xc of class c available inthe training dataset. We initialize with these probabilitydistributions the class-conditional latent vector generator,that is a random process that takes as input a classlabel c and returns as output a latent vector Zc drawnat random from Nc. During the adversarial training,the probability distributions Nc are considered invariantforcing the generator not to diverge from the initial classencoding in the latent space.

Adversarial training. During the adversarial train-ing, data flows in batches through the generator G andthe discriminator D and their weights are fine tuned tooptimize their loss functions. The discriminator classifiesan input image as belonging to one of the n problem-specific classes or as being fake. For each batch we supply,1/(n+ 1) of the total images are fake, i.e. we provide thebest possible balance for the fake class. The fake data isgenerated as output of G that takes as inputs latent vec-tors Zc extracted from the class-conditional latent vectorgenerator. In turn, the class-conditional latent vectorgenerator takes as input uniformly distributed class la-bels c, i.e. the fake images are uniformly distributedbetween the problem-specific classes. When training thediscriminator D we optimize the sparse categorical cross-

entropy loss function to match the class labels for realimages and the fake label for the generated ones.

For every batch learnt by the discriminator, a batchof the same size is learnt by the generator G. To thisend, a batch of conditional latent vectors Zc is drawnat random by applying a uniform distribution on thelabels c. These vectors are processed by the generatorand the output images are fed into the discriminator.The parameters in G are optimized to match the labelsselected by the discriminator with the labels c used togenerate the images.

5 Results

We validate the proposed methodology on a set of fourdatasets. We consider: MNIST [17], CIFAR-10 [20],Flowers [21], and GTSRB [1]. The former two datasetsare well known, Flowers is a small dataset including realphotos of five categories of flowers that we reshaped tothe resolution of 224x224, and GTSRB is a traffic signrecognition dataset. Details on these datasets are shownin Table 1. The first three datasets are balanced, GT-SRB is imbalanced. We force imbalance in the first threedatasets by selecting a class and dropping a significantamount of its instances from the training set. We repeatthis process for each class and train different generativemodels for each resulting imbalanced dataset. The fol-lowing results for each class are always obtained whentraining with that class as minority class and we referto the images left out from the training set as droppedimages. Since GTSRB is already imbalanced, we do notfurther imbalance it.

We compare the proposed BAGAN model with thestate-of-the-art ACGAN model [2]. To the best of ourknowledge, ACGAN is the only methodology presentedin literature so far to consider class conditioning to drawimages of a target class starting from a dataset includ-ing multiple classes (Section 3). Both BAGAN andACGAN are trained on the target datasets by usingmajority and minority classes jointly. We also consider asimple GAN approach that learns to draw the minority-class images by training only on that class. For a faircomparison, we limit the architecture changes betweenthe considered methodologies (BAGAN , ACGAN , andGAN ). The difference between BAGAN and ACGAN arethose described in this paper (i.e. discriminator outputtopology and autoencoder-based initialization). For thesimple GAN , we adjust the reference ACGAN discrimi-nator output to discriminate only between real and fakeimages, and we remove the class conditioning for thegenerator input (this GAN is trained only over imagesfrom the minority class). Figures 4, 5, and 6 show a qual-itative analysis of representative images generated forCIFAR-10 and for the three most and least representedclasses in GTSRB . For CIFAR-10 we show results only

4

Table 1: Target datasets’ information including resolution, number of classes, and per-class image distributionstatistics for the training set.

Training images per classDataset name Resolution Classes Min Median Mean Max

MNIST 28×28 10 5421 5936 6000 6742CIFAR-10 32×32 10 5000 5000 5000 5000Flowers 224×224 5 533 599 634 798GTSRB 64×64 43 210 600 911 2250

for minority-class images. For each class, 40% of thatclass images are dropped, generative models are trained,and randomly generated images are shown, Figure 4.For CIFAR-10 , the simple GAN collapses towards thegeneration of a single image example per class. To trainthis GAN we use only 3000 minority-class images (40%of the minority-class images are dropped and majorityclasses are not included in the training). Adversarialnetworks need many examples to learn drawing new im-ages [19] and in this case the simple GAN collapses.For ACGAN and BAGAN this issue is less relevant be-cause they can learn features from minority and majorityclasses jointly. To better understand the different be-havior of ACGAN and BAGAN , let us focus on theGTSRB dataset Figures 5 and 6. This dataset is orig-inally imbalanced and we train the generative modelswithout modifying it. For the majority classes, both AC-GAN and BAGAN return high-quality results, Figures5(c) and 5(b). Nonetheless, ACGAN fails in drawingimages for the minority classes and collapses towards thegeneration of a single example for each class, Figure 6(c).

In some cases ACGAN produces images that are notrepresentative of the desired class, e.g. the second rowin Figure 6(c) should be a warning sign whereas a speedlimit is drawn. BAGAN is never rewarded for drawinga realistic image if this does not represent the desiredclass. Thus, BAGAN does not exhibit this behavior.

5.1 Quantitative Assessment of Gener-ated Images

Since our goal is to leverage the generative model toaugment an imbalanced dataset by generating additionalminority-class images, we aim at the following goals:

a) Generated images must represent the desired class.

b) Generated images must not be repetitive.

c) Generated images must be different from the real onesalready available in the training set.

Missing to meet a) means that the generative model isnot capable to generate images that accurately represent

(a) Real image samples (b) BAGAN (c) ACGAN (d) Simple GAN

Figure 4: Five representative samples for each class (row) in the CIFAR-10 dataset. For each class, these samplesare obtained with generative models trained after dropping from the training set 40% of the images of that specificclass.

5

(a) Real image samples (b) BAGAN (c) ACGAN (d) Simple GAN

Figure 5: Five representative samples generated for the three most represented majority classes in the GT-SRB dataset.

(a) Real image samples (b) BAGAN (c) ACGAN (d) Simple GAN

Figure 6: Five representative samples generated for the three least represented minority classes in the GT-SRB dataset.

the target class and they look either as real examplesof other classes or they do not look real. Missing tomeet b) means that the generative model collapsed tothe generation of a single or few modes. Missing to meetc) means that we simply learnt to redraw the availabletraining images. We assess the quality of the generatedimages on the basis of these three goals.

Accuracy of the generated images. To verify thatthe images generated by the considered methodologiesare representative of the desired classes, we classify themby means of a deep learning model trained on the wholeoriginal dataset and we verify if the predicted classesmatch the target ones. In this work we use a ResNet-18 model [22]. Results are shown in Figure 7. Thesimple GAN returns the worst accuracy for generatedimages. The proposed BAGAN approach is generallybetter than the other approaches and generates imagesthat the ResNet-18 model can classify with the highestaccuracy. We observe again that a strong imbalance can

significantly deteriorate the quality of generated imageswith an accuracy that decreases as the percentage ofdropped images increases. This phenomenon is mostevident for ACGAN when targeting the MNIST dataset.

Variability of generated images. We measure sim-ilarity between two images by means of the structuralimage similarity SSIM [23]. This metric predicts humanperceptual similarity judgment, it returns one when thetwo images are identical and decreases as differences be-come more relevant. To verify that generated images arediverse, for each class we repeatedly generate a coupleof images and measure their similarity SSIM. Figure 8shows this diversity analysis for the considered datasetsaveraged over all classes. For MNIST , CIFAR-10 , andFlowers , we vary the percentage of minority-class imagesdropped within the set {40, 60, 80, 90, 95, 97.5}, whereasfor GTSRB we use the originally imbalanced dataset.We include in the analysis also a reference value thatis the average SSIM between real image couples of the

0 0.2 0.4 0.6 0.8

1

40 60 80 90 95 97.5

Acc

ura

cy

MNIST

40 60 80 90 95 97.5

Percentage of minority-class images dropped from the train set [%]

CIFAR-10

40 60 80 90 95 97.5

Flowers

0

GTSRB

GANACGANBAGAN

Figure 7: Accuracy of the images generated by the considered methodologies when varying the percentage ofminority-class images dropped before training the generative models. The accuracy is based on a ResNet-18 classifiertrained without dropping any image.

6

0 0.2 0.4 0.6 0.8

1

40 60 80 90 95 97.5

SS

IM c

ou

ple

sMNIST

40 60 80 90 95 97.5

Percentage of minority-class images dropped from the train set [%]

CIFAR-10

40 60 80 90 95 97.5

Flowers

0

GTSRB

GANACGANBAGAN

Reference

Figure 8: Structural similarity for generated image couples (SSIM couples, y axis) when varying the percentage ofimages dropped from the training set (x axis).

same class. When taking a random couple of real im-ages for CIFAR-10 or Flowers, these have so little incommon that the reference SSIM gets very close to zero.In general real images are always more variable thanthe generated ones (lower SSIM). Variability in imagesgenerated by the simple GAN approach is very little andsampled image couples have SSIM very close to one. Theproposed BAGAN methodology exhibit the best variabil-ity with respect to GAN and ACGANwith SSIM valuesclosest to the reference. For CIFAR-10 and Flowers,all methodologies deteriorate for strong imbalances withSSIM values that increase with the percentage of imagesdropped from the training set.Image diversity with respect to the training

set. To assess the variability of generated images withrespect to the ones already available in the training set.We compute the SSIM between generated images andtheir closest real neighbour. We compare this value withrespect to the image variability in the training set, i.e.the SSIM value between a real image and its closestreal neighbour. These SSIM values are very close toeach others meaning that there was no overfitting. Thisstatement holds for all the considered methodologies.In particular SSIM values of about 0.8, 0.25, 0.05, and0.5 are measured respectively for MNIST , CIFAR-10 ,Flowers, and GTSRB .

5.2 Quality of the Final Classification

We finally assess the accuracy of a deep-learning classifiertrained on an augmented dataset. For MNIST , CIFAR-

10 , and Flowers, for each class we: 1) select this classas minority class, 2) generate an imbalanced dataset bydropping a percentage of images for this class from thetraining set, 3) train the considered generative models,4) augment the imbalanced dataset to restore its balanceby means of the generative models, 5) train a ResNet-18 classifier for the augmented dataset, and 6) measurethe classifier accuracy for the minority class over the testset. Since GTSRB is already imbalanced, for this datasetwe skip steps 1) and 2). Augmentations obtained by thegenerative models are compared to the plain imbalanceddataset and to an horizontal mirroring augmentationapproach (mirror) where new minority-class images aregenerated by mirroring the ones available in the trainingset.

Accuracy results averaged over the different classes areshown in Figure 9. The proposed BAGAN methodologyreturns the best accuracy for GTSRB and most of thetime also for MNIST . These two datasets are character-ized by features sensible to the image orientation andthe mirroring approach as expected returns the worstaccuracy results because it disrupts these features. ForCIFAR-10 and Flowers the best accuracy is achievedby using the mirroring approach. Mirroring for thesedatasets does not disrupt any feature, qualitatively themirrored images are as good as the original ones. TheBAGAN approach still provides the best accuracy whencompared to ACGAN and GAN .

From this analysis we conclude that BAGAN is su-perior to other state-of-the-art adversarial generativenetworks when aiming at the generation of minority-

0.965 0.97

0.975 0.98

0.985 0.99

0.995

40 60 80 90 95 97.5

Acc

ura

cy

MNIST

0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8

40 60 80 90 95 97.5

Percentage of minority-class images dropped from the train set [%]

CIFAR-10

0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9

40 60 80 90 95 97.5

Flowers

0.925 0.93

0.935 0.94

0.945 0.95

0.955 0.96

0.965 0.97

0

GTSRB

GANACGANBAGAN

PlainMirror

Figure 9: Average accuracy for the minority-class achieved with a ResNet-18 classifier trained with the augmenteddataset whose balance is restored after dropping a percentage of minority-class images.

7

class images starting from an imbalanced dataset. Ad-ditionally we conclude that: when it is not easy to aug-ment a dataset with traditional techniques because oforientation-related features, BAGAN can be applied toimprove the final classification accuracy.

6 Conclusion

In this work we presented a methodology to restore thebalance of an imbalanced dataset by using generative ad-versarial networks. In the proposed BAGAN frameworkthe generator and the discriminator modules are initial-ized by means of an autoencoder to start the adversarialtraining from a good solution and to learn how differentclasses should be represented in the latent space.

We compared the proposed methodology against thestate of the art. Empirical results demonstrate thatBAGAN is superior to other generative adversarial net-works when aiming at the generation of high qualityimages starting with an imbalanced training set. Thisin turn results in a higher accuracy of deep-learningclassifiers trained over the augmented dataset where thebalance has been restored.

References

[1] J. Stallkamp, M. Schlipsing, J. Salmen, and C. Igel.Man vs. computer: Benchmarking machine learn-ing algorithms for traffic sign recognition. NeuralNetworks, (0):–, 2012.

[2] Augustus Odena, Christopher Olah, and JonathonShlens. Conditional image synthesis with auxiliaryclassifier GANs. In Doina Precup and Yee WhyeTeh, editors, Proceedings of the 34th InternationalConference on Machine Learning, volume 70 of Pro-ceedings of Machine Learning Research, pages 2642–2651, International Convention Centre, Sydney, Aus-tralia, 06–11 Aug 2017. PMLR.

[3] Anh Nguyen, Alexey Dosovitskiy, Jason Yosinski,Thomas Brox, and Jeff Clune. Synthesizing thepreferred inputs for neurons in neural networks viadeep generator networks. In D. D. Lee, M. Sugiyama,U. V. Luxburg, I. Guyon, and R. Garnett, editors,Advances in Neural Information Processing Systems29, pages 3387–3395. Curran Associates, Inc., 2016.

[4] A. Srivastava, L. Valkov, C. Russell, M. U. Gutmann,and C. Sutton. VEEGAN: Reducing Mode Collapsein GANs using Implicit Variational Learning. ArXive-prints, May 2017.

[5] Kevin Roth, Aurelien Lucchi, Sebastian Nowozin,and Thomas Hofmann. Stabilizing training of gen-erative adversarial networks through regularization.

In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach,R. Fergus, S. Vishwanathan, and R. Garnett, ed-itors, Advances in Neural Information ProcessingSystems 30, pages 2015–2025. Curran Associates,Inc., 2017.

[6] Naveen Kodali, Jacob D. Abernethy, James Hays,and Zsolt Kira. How to train your DRAGAN. CoRR,abs/1705.07215, 2017.

[7] M. Lucic, K. Kurach, M. Michalski, S. Gelly, andO. Bousquet. Are GANs Created Equal? A Large-Scale Study. ArXiv e-prints, #nov# 2017.

[8] I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu,D. Warde-Farley, S. Ozair, A. Courville, and Y. Ben-gio. Generative Adversarial Networks. ArXiv e-prints, June 2014.

[9] Alec Radford, Luke Metz, and Soumith Chintala.Unsupervised representation learning with deep con-volutional generative adversarial networks. CoRR,abs/1511.06434, 2015.

[10] Min Lin. Softmax GAN. CoRR, abs/1704.06191,2017.

[11] Luke Metz, Ben Poole, David Pfau, and Jascha Sohl-Dickstein. Unrolled generative adversarial networks.CoRR, abs/1611.02163, 2016.

[12] T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung,A. Radford, and X. Chen. Improved techniques fortraining gans. ArXiv e-prints, June 2016.

[13] A. Nguyen, J. Clune, Y. Bengio, A. Dosovitskiy,and J. Yosinski. Plug & play generative networks:Conditional iterative generation of images in latentspace. ArXiv e-prints, November 2016.

[14] V. Dumoulin, I. Belghazi, B. Poole, O. Mastropietro,A. Lamb, M. Arjovsky, and A. Courville. Adver-sarially Learned Inference. ArXiv e-prints, June2016.

[15] J. Donahue, P. Krahenbuhl, and T. Darrell. Ad-versarial Feature Learning. ArXiv e-prints, May2016.

[16] Luis Perez and Jason Wang. The effectiveness ofdata augmentation in image classification using deeplearning. CoRR, abs/1712.04621, 2017.

[17] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner.Gradient-based learning applied to document recog-nition. Proceedings of the IEEE, 86(11):2278–2324,Nov 1998.

[18] Tero Karras, Timo Aila, Samuli Laine, and JaakkoLehtinen. Progressive growing of GANs for improvedquality, stability, and variation. In InternationalConference on Learning Representations, 2018.

8

[19] Swaminathan Gurumurthy, Ravi Kiran Sarvadevab-hatla, and Venkatesh Babu Radhakrishnan. Deligan: Generative adversarial networks for diverse andlimited data. CoRR, abs/1706.02071, 2017.

[20] Alex Krizhevsky. Learning multiple layers of featuresfrom tiny images. Master’s thesis, 2009.

[21] Image classification using flowers dataset, 2018.https://cloud.google.com/ml-engine/docs/flowers-tutorial.

[22] K. He, X. Zhang, S. Ren, and J. Sun. Deep resid-ual learning for image recognition. In 2016 IEEE

Conference on Computer Vision and Pattern Recog-nition (CVPR), pages 770–778, June 2016.

[23] Zhou Wang, A. C. Bovik, H. R. Sheikh, and E. P.Simoncelli. Image quality assessment: from errorvisibility to structural similarity. IEEE Transactionson Image Processing, 13(4):600–612, April 2004.

[24] A. Dosovitskiy and T. Brox. Generating Imageswith Perceptual Similarity Metrics based on Deep

Networks. ArXiv e-prints, February 2016.

9