15
$\begingroup$

I'm reworking some of the GANs I originally made in TensorFlow2 to see if I can improve performance in Mathematica, and have been stuck on how to create a custom Minibatch Standard Deviation Layer. I'm trying to implement it to stabilize the training process and reduce instances of Mode Collapse. (More information on its purpose (with Python) can be found here.) Basically it takes the mean of the samples and figures out the standard deviation within each batch. The generator is penalized if the standard deviation is too low.

I thought the new FunctionLayer in 12.2 would be useful for this purpose but this is as far as I've gotten so far. Some how I need to thread that minibatch layer back into the network.

discriminator = NetGraph[{
   ConvolutionLayer[16, 5, "Stride" -> 3, "Input" -> {3, 4000}],
   ElementwiseLayer["ELU"],
   ConvolutionLayer[32, 3, "Stride" -> 1],
   ElementwiseLayer["ELU"],
   ConvolutionLayer[64, 2, "Stride" -> 1],
   ElementwiseLayer["ELU"],
   minibatch,
   FlattenLayer[],
   100,
   Ramp,
   1,
   Tanh,
   SequenceLastLayer[]
   },
  {1 -> 2 -> 3 -> 4 -> 5 -> 6 -> 8 -> 9 -> 10 -> 11 -> 12 -> 13, 
   6 -> 7}
]

minibatch = NetChain[{
   AggregationLayer[Mean, {2}],
   FunctionLayer[StandardDeviation]
   }]

GAN = NetInitialize@NetGANOperator[{generator, discriminator}]

BATCHSIZE = 64;
noise := RandomVariate[NormalDistribution[0, 0.2], 400]

trained = NetTrain[
  GAN,
  <|"Sample" -> RandomSample[trainingData, BATCHSIZE], 
   "Latent" -> Table[noise, BATCHSIZE]|>,
  BatchSize -> BATCHSIZE
  ]

enter image description here

The Wolfram framework always seems to apply the whole net to one example. How can I apply a function to all samples in the round and then append it to the array? Even if it's just a reference, I'm pretty sure I can follow it through.

Updated 2020-12-18 Along with the bounty I have also added a 3 samples of the training data that can be seen on a public gist of the shape {3,300}.

Updated 2021-01-02 Corrected code.

$\endgroup$
1
  • $\begingroup$ It looks like your code not complete. There is a message about it Syntax::bktmcp: Expression "NetGraph[<<1>>" has no closing "]". $\endgroup$ Jan 1, 2021 at 18:27

1 Answer 1

3
+200
$\begingroup$

I am not sure this is correct, but here's my attempt - perhaps it can get you started and/or inspire others to give a more correct answer.

My understanding is that we want access to the standard deviation of some features across the batches during training. BatchNormalizationLayer should have access to this internally - even though the only exposed ports are Input and Output. In-fact the Properties & Relations section under Documentation gives us a hint, that it actually computes:

batchNormFunction = 
  Function[Block[{sd = Sqrt[#MovingVariance + #Epsilon]},
    (#2 * #Scaling / sd ) + (#Biases - (#Scaling * #MovingMean)/sd)]];

As such, I think we can (ab)use the layer to extract the batch standard deviation as follows:

params = <|"Scaling" -> 1, "Biases" -> 0, "MovingMean" -> 0, 
   "MovingVariance" -> mv, "Epsilon" -> 0|>;
input/batchNormFunction[params, input]

Out[84]= Sqrt[mv]

Note this will fail if either "MovingVariance" or the input is zero. So you probably want to add a small amount of "Biases" and "Epsilon" to stabilize it. Here I define a 'frozen' BatchNormalizationLayer with no "Biases" and the default value of "Epsilon" (0.001) as follows:

bn = BatchNormalizationLayer["Scaling" -> 1, "Biases" -> 0, 
  "MovingMean" -> 0, "Input" -> {512, 4, 4}, 
  LearningRateMultipliers -> <|"Biases" -> 0, "MovingMean" -> 0, 
    "MovingVariance" -> 1, "Scaling" -> 0|>]

and define our minibatch layer using the NetGraph:

minibatch = 
 NetGraph[<|"bn" -> bn, "average" -> AggregationLayer[Mean, All], 
   "rep" -> ReplicateLayer[{1, 4, 4}], "cat" -> CatenateLayer[], 
   "divide" -> ThreadingLayer[Divide]|>, {NetPort["Input"] -> 
    "bn", {NetPort["Input"], "bn"} -> 
    "divide" -> "average" -> "rep", {NetPort["Input"], "rep"} -> 
    "cat"}]

enter image description here

Note that {512,4,4} above comes from the particular GAN I was playing around with, which I suspect will be easy enough to change. For reference here is the the GAN network - it seems to initialize properly, although I haven't trained on anything.

leaky = ParametricRampLayer["Slope" -> 0.2, 
   LearningRateMultipliers -> 0];
conv[channels_, kernel_] := 
 Sequence[ConvolutionLayer[channels, kernel, 
   PaddingSize -> (kernel + 1)/2 - 1], leaky]
deconv[channels_, kernel_] := 
 Sequence[DeconvolutionLayer[channels, kernel, 
   PaddingSize -> (kernel + 1)/2 - 1], leaky]
enc = NetEncoder[{"Image", {128, 128}}];
downSample[] := PoolingLayer[2, 2, "Function" -> Mean]
upSample[] := ResizeLayer[{Scaled[2], Scaled[2]}]

discriminator = 
 NetChain[{conv[16, 1], conv[16, 3], conv[32, 3], downSample[], 
   conv[32, 3], conv[64, 3], downSample[], conv[64, 3], conv[128, 3], 
   downSample[], conv[128, 3], conv[256, 3], downSample[], 
   conv[256, 3], conv[512, 3], downSample[], minibatch, conv[512, 3], 
   ConvolutionLayer[512, 4], leaky, LinearLayer[{}]}, "Input" -> enc]

generator = 
 NetChain[{DeconvolutionLayer[512, 4], leaky, deconv[512, 3], 
   upSample[], deconv[256, 3], deconv[256, 3], upSample[], 
   deconv[128, 3], deconv[128, 3], upSample[], deconv[64, 3], 
   deconv[64, 3], upSample[], deconv[32, 3], deconv[32, 3], 
   upSample[], deconv[16, 3], deconv[16, 3], 
   DeconvolutionLayer[3, 1]}, "Input" -> {512, 1, 1}]

NetInitialize[NetGANOperator[{generator, discriminator}]]

$\endgroup$
1
  • 2
    $\begingroup$ The timing is poor for me to try this. I've been away from the computer today and need to get in a plane first thing tomorrow. Reading through it looks like an excellent start. I appreciate the contrabution and will return to it soon.. $\endgroup$
    – BBirdsell
    Jan 3, 2021 at 4:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.