Generative Adversarial Network

Question

Generative adversarial networks (GAN) is regarded as one of "the most interesting idea in the last ten years in machine learning" by Yann LeCun. It can be used to generate photo-realistic images that are almost indistinguishable from the real ones.

GAN trains two competing neural networks: a generator network which generates image, and a discriminator network that distinguishes the generated image and the real training image. For example, the images shown below are generated by the network from the texts above them (taken from Han Zhang, etc., StackGAN: Text to Photo-realistic Image Synthesis with Stacked Generative Adversarial Networks).

I'm wondering whether we can implement a simplified version of that in Mathematica, given that the neural network framework has be enhanced greatly in version 11.1.

Answer 1

I wrestled with this for a while and got some kind of results, but nowhere near the great performance for which GANs are famous. Ultimately, they're absurdly sensitive to hyperparameters and initialization, and if you don't exactly imitate the published settings, you are unlikely to get good results.

I figure I should post my attempt -- maybe the community can figure out a good set of parameters that works. Mine sort of trained, but suffered from mode collapse and often converged on a blob. However, for sets of all one digit, it did seem to work okay, although this is much easier and not really where GANs have any advantage.

I tried to implement the Wasserstein GAN (paper available here) to generate MNIST digits. The training procedure is to update the discriminator on a batch 5 times for every 1 time you show it to the generator. Because Mathematica doesn't yet allow preservation of optimizer parameters between calls to NetTrain, I couldn't get this to work. Instead, I trained the networks jointly as suggested by Taliesin Beynon, setting the learning rate on the generator to something like -1/5, because it seemed like a plausible approximation.

The paper also used RMSProp as an optimizer. Mathematica has an RMSProp option, but on the net I defined it immediately diverged no matter what learning rate I chose. I used ADAM instead.

To begin, let's get a big batch of MNIST digits.

mnist = ResourceData["MNIST"];
mnistDigits = First /@ mnist;

Let's give a 10-dimensional noise input to the generator, and define the generator and discriminator. Notice that the discriminator does not have an activation on its output -- this is specific to the WGAN, a normal GAN would have a LogisticSigmoid or something.

randomDim = 10;
generator = 
 NetChain[{128, Ramp, 128, Ramp, 28*28, LogisticSigmoid, 
   ReshapeLayer[{1, 28, 28}]}, "Input" -> randomDim]
discriminator = 
 NetChain[{128, Ramp, 128, Ramp, 128, Ramp, 1}, 
  "Input" -> {1, 28, 28}]

Now the tricky part. We'll feed noise into the generator to produce a fake image, and also accept a real image as input. We want to apply the discriminator to both images, but with one set of weights, so we concatenate them and use NetMapOperator. Then, the loss function should be to maximize the score on the real image while minimizing the score on the fake image, so we negate the real score and then add them.

wganNet = 
 NetInitialize[
  NetGraph[<|"gen" -> generator, 
    "discrimop" -> NetMapOperator[discriminator], 
    "cat" -> CatenateLayer[], 
    "reshape" -> ReshapeLayer[{2, 1, 28, 28}], 
    "flat" -> FlattenLayer[], "total" -> SummationLayer[], 
    "scale" -> 
     ConstantTimesLayer["Scaling" -> {-1, 1}]|>, {NetPort["random"] ->
      "gen" -> "cat", NetPort["Input"] -> "cat", 
    "cat" -> 
     "reshape" -> "discrimop" -> "flat" -> "scale" -> "total"}, 
   "Input" -> {1, 28, 28}]]

One of the strengths of Mathematica's neural networks framework is that it's really easy to watch the networks train. We'll feed the trainer a progress function that takes 4 fixed random inputs and shows the generator's output, so we can watch the generator evolve over time.

ClearAll[progressFuncCreator]
progressFuncCreator[rands_List] := 
 Function[{reals}, 
    ImageResize[
     NetDecoder[{"Image", "Grayscale"}][
      NetExtract[#Net, "gen"][reals]], 50]] /@ rands &

Finally, create the training data:

trainingData = <|"random" -> RandomReal[{-1, 1}, {randomDim}], 
     "Input" -> ArrayReshape[ImageData[#], {1, 28, 28}]|> & /@ 
   mnistDigits;

And train, watching the generator make a bunch of vaguely number-shaped blobs. Notice the "WeightClipping" option on the discriminator -- this is the "secret sauce" in Wasserstein GANs that makes them learn an approximation of the Wasserstein/Earth-Mover's distance as opposed to the Jensen-Shannon distance, as explained in the paper.

NetTrain[wganNet, trainingData, "Output", 
 Method -> {"ADAM", "Beta1" -> 0.5, "LearningRate" -> 0.00005, 
   "WeightClipping" -> {"discrimop" -> 0.01}},
 TrainingProgressReporting -> 
  progressFuncCreator[Table[RandomReal[{-1, 1}, {randomDim}], 4]], 
 LearningRateMultipliers -> {"scale" -> 0, "gen" -> -0.2}, 
 TargetDevice -> "GPU", BatchSize -> 64]

Overall, my impression of the neural networks framework is very good. It's extremely flexible, coherently designed, and also extremely pretty. Crucially, it's easier to watch your net train than in any other framework. However, due to difficulties with staged training/saving optimizer parameters, it's not yet possible to replicate (in the sense of replicating a scientific experiment) some published results, like GANs, that use weirder architectures.

Answer 2

Yes it is possible. You can do alternating training manually by literally following the algorithm, so that you have a Do loop whose body contains two calls to NetTrain, but that suffers from overhead at each alternation (this could be overcome with clever caching, but we haven't done that yet). An approximation of this is to build a single network and optimize the D and G losses simultaneously by using a negative learning rate for the generator.

I have prototyped this, but only on a toy example.

I encourage you to try how to do it, it didn't take us more than a few hours of playing around to make a simple GAN in which the data distribution is a gaussian, the discriminator is an MLP, and the generator is a single EmbeddingLayer (just a fixed set of samples that can be moved around by gradient updates).

Answer 3

I see the generation of images from classification models as a trivial idea, largely, because this kind of processes and algorithms are standard in natural language processing.

More to the point of the question, the following posts show generation of images with classification derived bases:

• an answer to "Code that generates a mandala"

• "Comparison of dimension reduction algorithms over mandala images generation"

In some sense when using SVD or NNMF bases to recognize an image we reconstruct it by appropriate overlaying of basis images. Obviously such overlaying can be done without a recognition goal just to generate new images.

Update, 2017-06-24

Looking at the answer of Michael Curry and running the code (and Taliesin Beynon's code) I kind of see using Neural Networks as some sort of a long route. The MNIST based images those codes generate can be generated in a much quicker and controllable way using SVD and NNMF.

As an example, examine this basis images of handwriting "5" obtained with NNMF:

With those kind of bases are generated the (re)constructed handwritten digit images from this MSE answer:

(The linked answer describes the generation procedure.)

Answer 4

This is a GAN using more accurate loss function,Thanks for @Michael Curry

mnist = ResourceData["MNIST"];
mnistDigits = First /@ mnist;

randomDim = 10;
generator = NetChain[{128, Ramp, 128, Ramp, 28*28, LogisticSigmoid, 
                      ReshapeLayer[{1, 28, 28}]}, "Input" -> randomDim];
discriminator = NetChain[{128, Ramp, 128, Ramp, 128, Ramp, 
                          1, LogisticSigmoid}, 
                          "Input" -> {1, 28, 28}];

The loss function follow this paper

Gloss = NetChain[{PartLayer[1], ElementwiseLayer[Log[1 - #] &]}];

Dloss = NetGraph[{PartLayer[1], ElementwiseLayer[Log[1 - #] &], 
                  PartLayer[2], ElementwiseLayer[Log[#] &], 
                  TotalLayer[], 
                  ElementwiseLayer[-1 # &](*transfer Maximize to Minimize*)}, 
                  {1 -> 2, 3 -> 4, {2, 4} -> 5 -> 6}];

In the flat layer,the first part estimate the probability of fake,the second part estimate the probability of real.You can add another port Output2 to check this

ganNet = NetInitialize[
            NetGraph[<|"gen" -> generator, 
                       "discrimop" -> NetMapOperator[discriminator], 
                       "cat" -> CatenateLayer[], 
                       "reshape" -> ReshapeLayer[{2, 1, 28, 28}], 
                       "flat" -> FlattenLayer[], "G_loss" -> Gloss, 
                       "D_loss" -> Dloss|>, 
                       {NetPort["random"] -> "gen" -> "cat", NetPort["Input"] -> "cat", "cat" -> "reshape" -> "discrimop" -> "flat" -> {"G_loss", "D_loss"}, 
                       "G_loss" -> NetPort["Loss1"], "D_loss" -> NetPort["Loss2"]}, 
                       "Input" -> NetEncoder[{"Image", {28, 28}, ColorSpace -> "Grayscale"}]]]

trainingData = <|"random" -> RandomVariate[NormalDistribution[], {Length@mnistDigits, randomDim}], 
                 "Input" -> mnistDigits|>;

But when training,I want to reduce the loss of "Loss1" and "Loss2".But find nothing in docs,only got code like this trained = NetTrain[net, <|"Input" -> trainingImages|>, "Loss"].But here we have two loss function.

In the Pytorch,we can use two optimizer to reduce the two loss.But in MMA,I have no idea.

So this a simple way to train the net.

ganNet = NetTrain[ganNet, trainingData, {"Loss1", "Loss2"}, 
                   Method -> {"ADAM", "Beta1" -> 0.5, "LearningRate" -> 0.00005, "WeightClipping" -> {"discrimop" -> 0.01}}, 
                   BatchSize -> 64];

Answer 5

Please ignore the whole text below, just don't use MMA for nn learning NOW — you'll spend unbelievable amounts of time into studing MMA way of doing thins and step into many bugs on your ways.

If your goal is not to become an expert in Wolfram Mathematica NN framework — don't go this way. Invest your time in integration with python et al.

Talking about GANs in the NeuralNetworks settings... I'd say — do not use it.

Now we have NetGANOperator and it works IF you could reproduce the original loss function used via setting "Wasserstein" as a loss setting to the NetGANOperator.

It will take some creativity on your side and data jangling to make your discriminator and generator to fit into NetGANOperator assumptions (one input/one output etc).

If you don't believe me just look at all the "trained separately" sections in all the "Construction notebook" sections of the NetModel.

I greatly appreciate the effort all the people put into make NNs easy in MMA, but the result is what we have. I personally think that reproducing more low level functionality like the common "forward/reverse/get gradients" would be good.

Probably you'll spend much more time than it will take to adapt the original python source to your needs.

There's no way I'm aware of to reproduce GAN behaviour (without using NetGANOperator). Original suggestions from WRI to use NetMapOperator for discriminator, or use NetInsertSharedArrays just do not work, so they finally implemented NetGANOperator.

There're lots of subtleties like BatchNormalizationLayer. So you'll spend your time investigating NeuralNetwork framework (and bugs in the NN framework). That time could be spent in generating python code btw.

So just call Python if you have to use MMA for machine learning. NeuralNetworks framework is finely designed, the only problem with it — it's of little usability beyond the documentation examples (biased opinion here).

And you will NOT be able to easily reproduce what everybody does in frameworks like PyTorch, Tensorflow and so on.

WRI just decided they know better so they came with the whole NetTrain thing, but you're not allowed to pass gradients or change the graph etc etc. Just use what we provided to you.

The code is behind MXNetLink library — you can't study or modify it. If somebody decided you don't need some functionality — you're on your own.

So just adapt Python and what the world does and use ExternalSessionEvaluate. The idea of having everything inside MMA is great but just does not work in cases of rapidly developing knowledge areas.

I don't know how much creativity and stamina was put into NeuralNetworks (inference, declarative approach, visualisation, net surgery etc etc etc), but it now looks like the tool to run the examples. Buggy, btw (like NetEncoder["Image"] leaking memory for files and unusable as of late 2021).

Sorry, WRI people! you're the greatest!!!!!!

P.S. I finally demoed the classic Pix2Pix training. Because I love Mathematica:)