Refactor GAN Code and Implement Policies

src/mlpack/methods/ann/CMakeLists.txt

@@ -3,8 +3,9 @@
 set(SOURCES
   ffn.hpp
   ffn_impl.hpp
-  gan_impl.hpp
   gan.hpp
+  gan_impl.hpp
+  gan_policies.hpp
   rnn.hpp
   rnn_impl.hpp
 )

src/mlpack/methods/ann/ffn.hpp

@@ -403,7 +403,8 @@ class FFN
   template<
     typename Model,
     typename InitializerType,
-    class NoiseType
+    typename NoiseType,
+    typename PolicyType
   >
   friend class GAN;
 }; // class FFN

src/mlpack/methods/ann/gan.hpp

@@ -14,6 +14,7 @@
 #include <mlpack/core.hpp>
 
 #include <mlpack/methods/ann/ffn.hpp>
+#include <mlpack/methods/ann/gan_policies.hpp>
 #include <mlpack/methods/ann/visitor/output_parameter_visitor.hpp>
 #include <mlpack/methods/ann/visitor/reset_visitor.hpp>
 #include <mlpack/methods/ann/visitor/weight_size_visitor.hpp>
@@ -31,7 +32,7 @@ namespace ann /** Artificial Neural Network. **/ {
  * technique can generate photographs that look at least superficially
  * authentic to human observers, having many realistic characteristics.
  *
- * For more information, see the following papers.
+ * For more information, see the following paper:
  *
  * @code
  * @article{Goodfellow14,
@@ -44,26 +45,17 @@ namespace ann /** Artificial Neural Network. **/ {
  *   eprint    = {1406.2661},
  * }
  * @endcode
- * 
- * @code
- * @article{Salimans16,
- *   author    = {Tim Salimans, Ian Goodfellow, Wojciech Zaremba,
- *                Vicki Cheung, Alec Radford and Xi Chen},
- *   title     = {Improved Techniques for Training GANs},
- *   year      = {2016},
- *   url       = {http://arxiv.org/abs/1606.03498},
- *   eprint    = {1606.03498},
- * }
- * @endcode
  *
  * @tparam Model The class type of Generator and Discriminator.
  * @tparam InitializationRuleType Type of Initializer.
  * @tparam Noise The noise function to use.
+ * @tparam PolicyType The GAN variant to be used (GAN, DCGAN, WGAN or WGANGP).
  */
 template<
   typename Model,
   typename InitializationRuleType,
-  class Noise
+  typename Noise,
+  typename PolicyType = StandardGAN
 >
 class GAN
 {
@@ -74,11 +66,13 @@ class GAN
    * @param trainData The real data.
    * @param generator Generator network.
    * @param discriminator Discriminator network.
-   * @param batchSize BatchSize to be used for training.
-   * @param generatorUpdateStep Number of steps of Discriminator training
-   *                            before updating generator.
-   * @param preTrainSize Num of preTraining step of Discriminator.
+   * @param batchSize Batch size to be used for training.
+   * @param generatorUpdateStep Number of steps to train Discriminator
+   *                            before updating Generator.
+   * @param preTrainSize Number of pre-training steps of Discriminator.
    * @param multiplier Ratio of learning rate of Discriminator to the Generator.
+   * @param clippingParameter Weight range for enforcing Lipschitz constraint.
+   * @param lambda Parameter for setting the gradient penalty.
    */
   GAN(arma::mat& trainData,
       Model& generator,
@@ -89,7 +83,9 @@ class GAN
       size_t batchSize,
       size_t generatorUpdateStep,
       size_t preTrainSize,
-      double multiplier);
+      double multiplier,
+      double clippingParameter = 0.01,
+      double lambda = 10.0);
 
   // Reset function.
   void Reset();
@@ -99,29 +95,103 @@ class GAN
   void Train(OptimizerType& Optimizer);
 
   /**
-   * Evaluate function for the GAN gives the performance of the GAN on the
+   * Evaluate function for the Standard GAN and DCGAN.
+   * This function gives the performance of the Standard GAN or DCGAN on the
    * current input.
    *
    * @param parameters The parameters of the network.
    * @param i Index of the current input.
+   * @param batchSize Variable to store the present number of inputs.
    */
-  double Evaluate(const arma::mat& parameters,
-                  const size_t i,
-                  const size_t batchSize);
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
+                          std::is_same<Policy, DCGAN>::value, double>::type
+  Evaluate(const arma::mat& parameters,
+           const size_t i,
+           const size_t batchSize);
 
   /**
-   * Gradient function for GAN.
+   * Evaluate function for the WGAN.
+   * This function gives the performance of the WGAN on the current input.
+   *
+   * @param parameters The parameters of the network.
+   * @param i Index of the current input.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGAN>::value,
+                          double>::type
+  Evaluate(const arma::mat& parameters,
+           const size_t i,
+           const size_t batchSize);
+
+  /**
+   * Evaluate function for the WGAN-GP.
+   * This function gives the performance of the WGAN-GP on the current input.
+   *
+   * @param parameters The parameters of the network.
+   * @param i Index of the current input.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                          double>::type
+  Evaluate(const arma::mat& parameters,
+           const size_t i,
+           const size_t batchSize);
+
+  /**
+   * Gradient function for Standard GAN and DCGAN.
    * This function passes the gradient based on which network is being
    * trained, i.e., Generator or Discriminator.
    * 
    * @param parameters present parameters of the network.
    * @param i Index of the predictors.
    * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
    */
-  void Gradient(const arma::mat& parameters,
-                const size_t i,
-                arma::mat& gradient,
-                const size_t batchSize);
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
+                          std::is_same<Policy, DCGAN>::value, void>::type
+  Gradient(const arma::mat& parameters,
+           const size_t i,
+           arma::mat& gradient,
+           const size_t batchSize);
+
+  /**
+   * Gradient function for WGAN.
+   * This function passes the gradient based on which network is being
+   * trained, i.e., Generator or Discriminator.
+   * 
+   * @param parameters present parameters of the network.
+   * @param i Index of the predictors.
+   * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGAN>::value, void>::type
+  Gradient(const arma::mat& parameters,
+           const size_t i,
+           arma::mat& gradient,
+           const size_t batchSize);
+
+  /**
+   * Gradient function for WGAN-GP.
+   * This function passes the gradient based on which network is being
+   * trained, i.e., Generator or Discriminator.
+   * 
+   * @param parameters present parameters of the network.
+   * @param i Index of the predictors.
+   * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                          void>::type
+  Gradient(const arma::mat& parameters,
+           const size_t i,
+           arma::mat& gradient,
+           const size_t batchSize);
 
   /**
    * Shuffle the order of function visitation. This may be called by the
@@ -167,7 +237,7 @@ class GAN
   //! Locally stored Discriminator network.
   Model& discriminator;
   //! Locally stored Initializer.
-  InitializationRuleType  initializeRule;
+  InitializationRuleType initializeRule;
   //! Locally stored Noise function
   Noise noiseFunction;
   //! Locally stored input dimension of the Generator network.
@@ -186,6 +256,10 @@ class GAN
   size_t preTrainSize;
   //! Locally stored learning rate ratio for Generator network.
   double multiplier;
+  //! Locally stored weight clipping parameter.
+  double clippingParameter;
+  //! Locally stored lambda parameter.
+  double lambda;
   //! Locally stored reset parameter.
   bool reset;
   //! Locally stored delta visitor.
@@ -208,6 +282,8 @@ class GAN
   arma::mat gradientDiscriminator;
   //! Locally stored gradient for noise data in the predictors.
   arma::mat noiseGradientDiscriminator;
+  //! Locally stored norm of the gradient of Discriminator.
+  arma::mat normGradientDiscriminator;
   //! Locally stored noise using the noise function.
   arma::mat noise;
   //! Locally stored gradient for Generator.

src/mlpack/methods/ann/gan_impl.hpp

@@ -23,8 +23,13 @@
 
 namespace mlpack {
 namespace ann /** Artifical Neural Network.  */ {
-template<typename Model, typename InitializationRuleType, class Noise>
-GAN<Model, InitializationRuleType, Noise>::GAN(
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
     arma::mat& predictors,
     Model& generator,
     Model& discriminator,
@@ -34,7 +39,9 @@ GAN<Model, InitializationRuleType, Noise>::GAN(
     size_t batchSize,
     size_t generatorUpdateStep,
     size_t preTrainSize,
-    double multiplier):
+    double multiplier,
+    double clippingParameter,
+    double lambda):
     predictors(predictors),
     generator(generator),
     discriminator(discriminator),
@@ -45,6 +52,8 @@ GAN<Model, InitializationRuleType, Noise>::GAN(
     generatorUpdateStep(generatorUpdateStep),
     preTrainSize(preTrainSize),
     multiplier(multiplier),
+    clippingParameter(clippingParameter),
+    lambda(lambda),
     reset(false)
 {
   // Insert IdentityLayer for joining the Generator and Discriminator.
@@ -77,8 +86,13 @@ GAN<Model, InitializationRuleType, Noise>::GAN(
   generator.responses.set_size(predictors.n_rows, batchSize);
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-void GAN<Model, InitializationRuleType, Noise>::Reset()
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::Reset()
 {
   size_t genWeights = 0;
   size_t discWeights = 0;
@@ -110,9 +124,14 @@ void GAN<Model, InitializationRuleType, Noise>::Reset()
   reset = true;
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
 template<typename OptimizerType>
-void GAN<Model, InitializationRuleType, Noise>::Train(
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::Train(
     OptimizerType& Optimizer)
 {
   if (!reset)
@@ -120,11 +139,19 @@ void GAN<Model, InitializationRuleType, Noise>::Train(
   Optimizer.Optimize(*this, parameter);
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-double GAN<Model, InitializationRuleType, Noise>::Evaluate(
-    const arma::mat& /*parameters*/,
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
+                        std::is_same<Policy, DCGAN>::value, double>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
+    const arma::mat& /* parameters */,
     const size_t i,
-    const size_t /*batchSize*/)
+    const size_t /* batchSize */)
 {
   if (!reset)
     Reset();
@@ -137,8 +164,8 @@ double GAN<Model, InitializationRuleType, Noise>::Evaluate(
   discriminator.Forward(std::move(currentInput));
   double res = discriminator.outputLayer.Forward(
       std::move(boost::apply_visitor(
-          outputParameterVisitor,
-          discriminator.network.back())), std::move(currentTarget));
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
 
   noise.imbue( [&]() { return noiseFunction();} );
   generator.Forward(std::move(noise));
@@ -153,17 +180,134 @@ double GAN<Model, InitializationRuleType, Noise>::Evaluate(
   currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
       1, batchSize, false, false);
   res += discriminator.outputLayer.Forward(
-        std::move(boost::apply_visitor(
-            outputParameterVisitor,
-            discriminator.network.back())), std::move(currentTarget));
+      std::move(boost::apply_visitor(
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
 
   return res;
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-void GAN<Model, InitializationRuleType, Noise>::
-Gradient(const arma::mat& /*parameters*/, const size_t i, arma::mat& gradient,
-    const size_t /*batchSize*/)
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, WGAN>::value, double>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
+    const arma::mat& /* parameters */,
+    const size_t i,
+    const size_t /* batchSize */)
+{
+  if (!reset)
+    Reset();
+
+  currentInput = arma::mat(predictors.memptr() + (i * predictors.n_rows),
+      predictors.n_rows, batchSize, false, false);
+  currentTarget = arma::mat(responses.memptr() + i, 1, batchSize, false,
+      false);
+
+  discriminator.Forward(std::move(currentInput));
+  double res = discriminator.outputLayer.Forward(
+      std::move(boost::apply_visitor(
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
+
+  noise.imbue( [&]() { return noiseFunction();} );
+  generator.Forward(std::move(noise));
+
+  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+      boost::apply_visitor(outputParameterVisitor, generator.network.back());
+  discriminator.Forward(std::move(discriminator.predictors.cols(numFunctions,
+      numFunctions + batchSize - 1)));
+  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+      -arma::ones(1, batchSize);
+
+  currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
+      1, batchSize, false, false);
+  res += discriminator.outputLayer.Forward(
+      std::move(boost::apply_visitor(
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
+
+  return res;
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                        double>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
+    const arma::mat& /* parameters */,
+    const size_t i,
+    const size_t /* batchSize */)
+{
+  if (!reset)
+    Reset();
+
+  currentInput = arma::mat(predictors.memptr() + (i * predictors.n_rows),
+      predictors.n_rows, batchSize, false, false);
+  currentTarget = arma::mat(responses.memptr() + i, 1, batchSize, false,
+      false);
+
+  discriminator.Forward(std::move(currentInput));
+  double res = discriminator.outputLayer.Forward(
+      std::move(boost::apply_visitor(
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
+
+  noise.imbue( [&]() { return noiseFunction();} );
+  generator.Forward(std::move(noise));
+
+  arma::mat generatedData = boost::apply_visitor(outputParameterVisitor,
+      generator.network.back());
+  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+      generatedData;
+  discriminator.Forward(std::move(discriminator.predictors.cols(numFunctions,
+      numFunctions + batchSize - 1)));
+  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+      -arma::ones(1, batchSize);
+
+  currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
+      1, batchSize, false, false);
+  res += discriminator.outputLayer.Forward(
+      std::move(boost::apply_visitor(
+      outputParameterVisitor,
+      discriminator.network.back())), std::move(currentTarget));
+
+  // Gradient Penalty is calculated here.
+  double epsilon = math::Random();
+  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+      (epsilon * currentInput) + ((1.0 - epsilon) * generatedData);
+  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+      -arma::ones(1, batchSize);
+  discriminator.Gradient(discriminator.parameter, numFunctions,
+      normGradientDiscriminator, batchSize);
+  res += lambda * std::pow(arma::norm(normGradientDiscriminator, 2) - 1, 2);
+
+  return res;
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
+                        std::is_same<Policy, DCGAN>::value, void>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::
+Gradient(const arma::mat& /* parameters */,
+         const size_t i,
+         arma::mat& gradient,
+         const size_t /* batchSize */)
 {
   if (!reset)
     Reset();
@@ -215,7 +359,7 @@ Gradient(const arma::mat& /*parameters*/, const size_t i, arma::mat& gradient,
     discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
         arma::ones(1, batchSize);
     discriminator.Gradient(discriminator.parameter, numFunctions,
-      noiseGradientDiscriminator, batchSize);
+        noiseGradientDiscriminator, batchSize);
     generator.error = boost::apply_visitor(deltaVisitor,
         discriminator.network[1]);
 
@@ -227,30 +371,222 @@ Gradient(const arma::mat& /*parameters*/, const size_t i, arma::mat& gradient,
   }
 
   counter++;
+  currentBatch++;
 
-  if (counter >= numFunctions)
+  // Revert the counter to zero, if the total dataset get's covered.
+  if (counter * batchSize >= numFunctions)
   {
     counter = 0;
-    currentBatch++;
   }
-  else if (counter % batchSize == 0)
+
+  if (preTrainSize > 0)
   {
-    currentBatch++;
-    if (preTrainSize > 0)
-    {
-      preTrainSize--;
-    }
+    preTrainSize--;
   }
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-void GAN<Model, InitializationRuleType, Noise>::Shuffle()
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, WGAN>::value, void>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::
+Gradient(const arma::mat& /* parameters */,
+         const size_t i,
+         arma::mat& gradient,
+         const size_t /* batchSize */)
+{
+  if (!reset)
+    Reset();
+
+  if (gradient.is_empty())
+  {
+    if (parameter.is_empty())
+      Reset();
+    gradient = arma::zeros<arma::mat>(parameter.n_elem, 1);
+  }
+  else
+    gradient.zeros();
+
+  if (noiseGradientDiscriminator.is_empty())
+  {
+    noiseGradientDiscriminator = arma::zeros<arma::mat>(
+        gradientDiscriminator.n_elem, 1);
+  }
+  else
+  {
+    noiseGradientDiscriminator.zeros();
+  }
+
+  gradientGenerator = arma::mat(gradient.memptr(),
+      generator.Parameters().n_elem, 1, false, false);
+
+  gradientDiscriminator = arma::mat(gradient.memptr() +
+      gradientGenerator.n_elem,
+      discriminator.Parameters().n_elem, 1, false, false);
+
+  // Get the gradients of the Discriminator.
+  discriminator.Gradient(discriminator.parameter, i, gradientDiscriminator,
+      batchSize);
+  noise.imbue( [&]() { return noiseFunction();} );
+  generator.Forward(std::move(noise));
+  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+      boost::apply_visitor(outputParameterVisitor, generator.network.back());
+
+  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+      -arma::ones(1, batchSize);
+  discriminator.Gradient(discriminator.parameter, numFunctions,
+      noiseGradientDiscriminator, batchSize);
+  gradientDiscriminator += noiseGradientDiscriminator;
+  gradientDiscriminator = arma::clamp(gradientDiscriminator,
+      -clippingParameter, clippingParameter);
+
+  if (currentBatch % generatorUpdateStep == 0 && preTrainSize == 0)
+  {
+    // Minimize -D(G(noise)).
+    // Pass the error from Discriminator to Generator.
+    discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+        arma::ones(1, batchSize);
+    discriminator.Gradient(discriminator.parameter, numFunctions,
+        noiseGradientDiscriminator, batchSize);
+    generator.error = boost::apply_visitor(deltaVisitor,
+        discriminator.network[1]);
+
+    generator.Predictors() = noise;
+    generator.ResetGradients(gradientGenerator);
+    generator.Gradient(generator.parameter, 0, gradientGenerator, batchSize);
+
+    gradientGenerator *= multiplier;
+  }
+
+  counter++;
+  currentBatch++;
+
+  // Revert the counter to zero, if the total dataset get's covered.
+  if (counter * batchSize >= numFunctions)
+  {
+    counter = 0;
+  }
+
+  if (preTrainSize > 0)
+  {
+    preTrainSize--;
+  }
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+template<typename Policy>
+typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                        void>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::
+Gradient(const arma::mat& /* parameters */,
+         const size_t i,
+         arma::mat& gradient,
+         const size_t /* batchSize */)
+{
+  if (!reset)
+    Reset();
+
+  if (gradient.is_empty())
+  {
+    if (parameter.is_empty())
+      Reset();
+    gradient = arma::zeros<arma::mat>(parameter.n_elem, 1);
+  }
+  else
+    gradient.zeros();
+
+  if (noiseGradientDiscriminator.is_empty())
+  {
+    noiseGradientDiscriminator = arma::zeros<arma::mat>(
+        gradientDiscriminator.n_elem, 1);
+  }
+  else
+  {
+    noiseGradientDiscriminator.zeros();
+  }
+
+  gradientGenerator = arma::mat(gradient.memptr(),
+      generator.Parameters().n_elem, 1, false, false);
+
+  gradientDiscriminator = arma::mat(gradient.memptr() +
+      gradientGenerator.n_elem,
+      discriminator.Parameters().n_elem, 1, false, false);
+
+  // Get the gradients of the Discriminator.
+  discriminator.Gradient(discriminator.parameter, i, gradientDiscriminator,
+      batchSize);
+  noise.imbue( [&]() { return noiseFunction();} );
+  generator.Forward(std::move(noise));
+  discriminator.predictors.cols(numFunctions, numFunctions + batchSize - 1) =
+      boost::apply_visitor(outputParameterVisitor, generator.network.back());
+
+  discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+      -arma::ones(1, batchSize);
+  discriminator.Gradient(discriminator.parameter, numFunctions,
+      noiseGradientDiscriminator, batchSize);
+  gradientDiscriminator += noiseGradientDiscriminator;
+
+  if (currentBatch % generatorUpdateStep == 0 && preTrainSize == 0)
+  {
+    // Minimize -D(G(noise)).
+    // Pass the error from Discriminator to Generator.
+    discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
+        arma::ones(1, batchSize);
+    discriminator.Gradient(discriminator.parameter, numFunctions,
+        noiseGradientDiscriminator, batchSize);
+    generator.error = boost::apply_visitor(deltaVisitor,
+        discriminator.network[1]);
+
+    generator.Predictors() = noise;
+    generator.ResetGradients(gradientGenerator);
+    generator.Gradient(generator.parameter, 0, gradientGenerator, batchSize);
+
+    gradientGenerator *= multiplier;
+  }
+
+  counter++;
+  currentBatch++;
+
+  // Revert the counter to zero, if the total dataset get's covered.
+  if (counter * batchSize >= numFunctions)
+  {
+    counter = 0;
+  }
+
+  if (preTrainSize > 0)
+  {
+    preTrainSize--;
+  }
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::Shuffle()
 {
   math::ShuffleData(predictors, responses, predictors, responses);
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-void GAN<Model, InitializationRuleType, Noise>::Forward(arma::mat&& input)
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::Forward(
+    arma::mat&& input)
 {
   if (!reset)
     Reset();
@@ -263,8 +599,13 @@ void GAN<Model, InitializationRuleType, Noise>::Forward(arma::mat&& input)
   discriminator.Forward(std::move(ganOutput));
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
-void GAN<Model, InitializationRuleType, Noise>::
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::
 Predict(arma::mat&& input, arma::mat& output)
 {
   if (!reset)
@@ -276,9 +617,14 @@ Predict(arma::mat&& input, arma::mat& output)
       discriminator.network.back());
 }
 
-template<typename Model, typename InitializationRuleType, typename Noise>
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
 template<typename Archive>
-void GAN<Model, InitializationRuleType, Noise>::
+void GAN<Model, InitializationRuleType, Noise, PolicyType>::
 serialize(Archive& ar, const unsigned int /* version */)
 {
   ar & BOOST_SERIALIZATION_NVP(parameter);

src/mlpack/methods/ann/gan_policies.hpp

@@ -0,0 +1,84 @@
+/**
+ * @file gan_policies.hpp
+ * @author Shikhar Jaiswal
+ *
+ * Implementation of the GAN policy types.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_GAN_POLICIES_HPP
+#define MLPACK_METHODS_ANN_GAN_POLICIES_HPP
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+/**
+ * For more information, see the following paper:
+ *
+ * @code
+ * @article{Salimans16,
+ *   author    = {Tim Salimans, Ian Goodfellow, Wojciech Zaremba,
+ *                Vicki Cheung, Alec Radford and Xi Chen},
+ *   title     = {Improved Techniques for Training GANs},
+ *   year      = {2016},
+ *   url       = {http://arxiv.org/abs/1606.03498},
+ *   eprint    = {1606.03498},
+ * }
+ * @endcode
+ */
+class StandardGAN { /* Nothing to do here */ };
+
+/**
+ * For more information, see the following paper:
+ *
+ * @code
+ * @article{Radford15,
+ *   author    = {Alec Radford, Luke Metz and Soumith Chintala},
+ *   title     = {Unsupervised Representation Learning with Deep Convolutional
+                  Generative Adversarial Networks},
+ *   year      = {2015},
+ *   url       = {https://arxiv.org/abs/1511.06434},
+ *   eprint    = {1511.06434},
+ * }
+ * @endcode
+ */
+class DCGAN { /* Nothing to do here */ };
+
+/**
+ * For more information, see the following paper:
+ *
+ * @code
+ * @article{Arjovsky17,
+ *   author    = {Martin Arjovsky, Soumith Chintala and Léon Bottou},
+ *   title     = {Wasserstein GAN},
+ *   year      = {2017},
+ *   url       = {https://arxiv.org/abs/1701.07875},
+ *   eprint    = {1701.07875},
+ * }
+ * @endcode
+ */
+class WGAN { /* Nothing to do here */ };
+
+/**
+ * For more information, see the following paper:
+ *
+ * @code
+ * @article{Gulrajani17,
+ *   author    = {Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent
+                  Dumoulin and Aaron Courville},
+ *   title     = {Improved Training of Wasserstein GANs},
+ *   year      = {2017},
+ *   url       = {https://arxiv.org/abs/1704.00028},
+ *   eprint    = {1704.00028},
+ * }
+ * @endcode
+ */
+class WGANGP { /* Nothing to do here */ };
+
+} // namespace ann
+} // namespace mlpack
+
+#endif

src/mlpack/methods/ann/loss_functions/CMakeLists.txt

@@ -3,6 +3,8 @@
 set(SOURCES
   cross_entropy_error.hpp
   cross_entropy_error_impl.hpp
+  earth_mover_distance.hpp
+  earth_mover_distance_impl.hpp
   kl_divergence.hpp
   kl_divergence_impl.hpp
   mean_squared_error.hpp

src/mlpack/methods/ann/loss_functions/cross_entropy_error.hpp

@@ -50,6 +50,7 @@ class CrossEntropyError
    */
   template<typename InputType, typename TargetType>
   double Forward(const InputType&& input, const TargetType&& target);
+
   /**
    * Ordinary feed backward pass of a neural network.
    *

src/mlpack/methods/ann/loss_functions/earth_mover_distance.hpp

@@ -0,0 +1,84 @@
+/**
+ * @file earth_mover_distance.hpp
+ * @author Shikhar Jaiswal
+ *
+ * Definition of the earth mover distance function.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_LOSS_FUNCTIONS_EARTH_MOVER_DISTANCE_HPP
+#define MLPACK_METHODS_ANN_LOSS_FUNCTIONS_EARTH_MOVER_DISTANCE_HPP
+
+#include <mlpack/prereqs.hpp>
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+/**
+ * The earth mover distance function measures the network's performance
+ * according to the Kantorovich-Rubinstein duality approximation.
+ *
+ * @tparam InputDataType Type of the input data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ * @tparam OutputDataType Type of the output data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ */
+template <
+    typename InputDataType = arma::mat,
+    typename OutputDataType = arma::mat
+>
+class EarthMoverDistance
+{
+ public:
+  /**
+   * Create the EarthMoverDistance object.
+   */
+  EarthMoverDistance();
+
+  /*
+   * Ordinary feed forward pass of a neural network.
+   *
+   * @param input Input data used for evaluating the specified function.
+   * @param output Resulting output activation.
+   */
+  template<typename InputType, typename TargetType>
+  double Forward(const InputType&& input, const TargetType&& target);
+
+  /**
+   * Ordinary feed backward pass of a neural network.
+   *
+   * @param input The propagated input activation.
+   * @param target The target vector.
+   * @param output The calculated error.
+   */
+  template<typename InputType, typename TargetType, typename OutputType>
+  void Backward(const InputType&& input,
+                const TargetType&& target,
+                OutputType&& output);
+
+  //! Get the output parameter.
+  OutputDataType& OutputParameter() const { return outputParameter; }
+  //! Modify the output parameter.
+  OutputDataType& OutputParameter() { return outputParameter; }
+
+  /**
+   * Serialize the layer.
+   */
+  template<typename Archive>
+  void serialize(Archive& ar, const unsigned int /* version */);
+
+ private:
+  //! Locally-stored output parameter object.
+  OutputDataType outputParameter;
+}; // class EarthMoverDistance
+
+} // namespace ann
+} // namespace mlpack
+
+// Include implementation.
+#include "earth_mover_distance_impl.hpp"
+
+#endif

src/mlpack/methods/ann/loss_functions/earth_mover_distance_impl.hpp

@@ -0,0 +1,57 @@
+/**
+ * @file earth_mover_distance_impl.hpp
+ * @author Shikhar Jaiswal
+ *
+ * Implementation of the earth mover distance function.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_LOSS_FUNCTIONS_EARTH_MOVER_DISTANCE_IMPL_HPP
+#define MLPACK_METHODS_ANN_LOSS_FUNCTIONS_EARTH_MOVER_DISTANCE_IMPL_HPP
+
+// In case it hasn't yet been included.
+#include "earth_mover_distance.hpp"
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+template<typename InputDataType, typename OutputDataType>
+EarthMoverDistance<InputDataType, OutputDataType>::EarthMoverDistance()
+{
+  // Nothing to do here.
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename InputType, typename TargetType>
+double EarthMoverDistance<InputDataType, OutputDataType>::Forward(
+    const InputType&& input, const TargetType&& target)
+{
+  return -arma::accu(target % input);
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename InputType, typename TargetType, typename OutputType>
+void EarthMoverDistance<InputDataType, OutputDataType>::Backward(
+    const InputType&& /* input */,
+    const TargetType&& target,
+    OutputType&& output)
+{
+  output = -target;
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename Archive>
+void EarthMoverDistance<InputDataType, OutputDataType>::serialize(
+    Archive& ar,
+    const unsigned int /* version */)
+{
+  /* Nothing to do here */
+}
+
+} // namespace ann
+} // namespace mlpack
+
+#endif

src/mlpack/tests/CMakeLists.txt

@@ -134,6 +134,7 @@ add_executable(mlpack_test
   ub_tree_test.cpp
   union_find_test.cpp
   vantage_point_tree_test.cpp
+  wgan_test.cpp
   main_tests/test_helper.hpp
   main_tests/emst_test.cpp
   main_tests/adaboost_test.cpp

src/mlpack/tests/dcgan_test.cpp

@@ -1,5 +1,5 @@
 /**
- * @file dcgan_network_test.cpp
+ * @file dcgan_test.cpp
  * @author Shikhar Jaiswal
  *
  * Tests the DCGAN network.
@@ -13,7 +13,6 @@
 
 #include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
 #include <mlpack/methods/ann/loss_functions/cross_entropy_error.hpp>
-#include <mlpack/methods/ann/loss_functions/sigmoid_cross_entropy_error.hpp>
 #include <mlpack/methods/ann/gan.hpp>
 #include <mlpack/methods/ann/ffn.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
@@ -67,8 +66,7 @@ BOOST_AUTO_TEST_CASE(DCGANMNISTTest)
   trainData.load("mnist_first250_training_4s_and_9s.arm");
   Log::Info << arma::size(trainData) << std::endl;
 
-  if (datasetMaxCols > 0)
-    trainData = trainData.cols(0, datasetMaxCols - 1);
+  trainData = trainData.cols(0, datasetMaxCols - 1);
 
   size_t numIterations = trainData.n_cols * numEpoches;
   numIterations /= batchSize;
@@ -78,7 +76,7 @@ BOOST_AUTO_TEST_CASE(DCGANMNISTTest)
   Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;
 
   // Create the Discriminator network
-  FFN<SigmoidCrossEntropyError<> > discriminator;
+  FFN<CrossEntropyError<> > discriminator;
   discriminator.Add<Convolution<> >(1, dNumKernels, 4, 4, 2, 2, 1, 1, 28, 28);
   discriminator.Add<LeakyReLU<> >(0.2);
   discriminator.Add<Convolution<> >(dNumKernels, 2 * dNumKernels, 4, 4, 2, 2,
@@ -95,36 +93,40 @@ BOOST_AUTO_TEST_CASE(DCGANMNISTTest)
   discriminator.Add<SigmoidLayer<> >();
 
   // Create the Generator network
-  FFN<SigmoidCrossEntropyError<> > generator;
+  FFN<CrossEntropyError<> > generator;
   generator.Add<TransposedConvolution<> >(noiseDim, 8 * dNumKernels, 2, 2,
       1, 1, 1, 1, 1, 1);
+  generator.Add<BatchNorm<> >(1024);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(8 * dNumKernels, 4 * dNumKernels,
       2, 2, 1, 1, 0, 0, 2, 2);
+  generator.Add<BatchNorm<> >(1152);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(4 * dNumKernels, 2 * dNumKernels,
       5, 5, 2, 2, 1, 1, 3, 3);
+  generator.Add<BatchNorm<> >(3136);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(2 * dNumKernels, dNumKernels, 8, 8,
       1, 1, 1, 1, 7, 7);
+  generator.Add<BatchNorm<> >(6272);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(dNumKernels, 1, 15, 15, 1, 1, 1, 1,
       14, 14);
   generator.Add<TanHLayer<> >();
 
-  // Create GAN
+  // Create DCGAN
   GaussianInitialization gaussian(0, 1);
   Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
       tolerance, shuffle);
   std::function<double()> noiseFunction = [] () {
       return math::RandNormal(0, 1);};
-  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
-      std::function<double()> > gan(trainData, generator, discriminator,
+  GAN<FFN<CrossEntropyError<> >, GaussianInitialization,
+      std::function<double()>, DCGAN> dcgan(trainData, generator, discriminator,
       gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
       discriminatorPreTrain, multiplier);
 
   Log::Info << "Training..." << std::endl;
-  gan.Train(optimizer);
+  dcgan.Train(optimizer);
 
   // Generate samples
   Log::Info << "Sampling..." << std::endl;
@@ -200,7 +202,7 @@ BOOST_AUTO_TEST_CASE(DCGANCelebATest)
   Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;
 
   // Create the Discriminator network
-  FFN<SigmoidCrossEntropyError<> > discriminator;
+  FFN<CrossEntropyError<> > discriminator;
   discriminator.Add<Convolution<> >(3, dNumKernels, 4, 4, 2, 2, 1, 1, 64, 64);
   discriminator.Add<LeakyReLU<> >(0.2);
   discriminator.Add<Convolution<> >(dNumKernels, 2 * dNumKernels, 4, 4, 2, 2,
@@ -217,36 +219,40 @@ BOOST_AUTO_TEST_CASE(DCGANCelebATest)
   discriminator.Add<SigmoidLayer<> >();
 
   // Create the Generator network
-  FFN<SigmoidCrossEntropyError<> > generator;
+  FFN<CrossEntropyError<> > generator;
   generator.Add<TransposedConvolution<> >(noiseDim, 8 * dNumKernels, 4, 4,
       1, 1, 2, 2, 1, 1);
+  generator.Add<BatchNorm<> >(4096);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(8 * dNumKernels, 4 * dNumKernels,
       5, 5, 1, 1, 1, 1, 4, 4);
+  generator.Add<BatchNorm<> >(8192);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(4 * dNumKernels, 2 * dNumKernels,
       9, 9, 1, 1, 1, 1, 8, 8);
+  generator.Add<BatchNorm<> >(16384);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(2 * dNumKernels, dNumKernels, 17, 17,
       1, 1, 1, 1, 16, 16);
+  generator.Add<BatchNorm<> >(32768);
   generator.Add<ReLULayer<> >();
   generator.Add<TransposedConvolution<> >(dNumKernels, 3, 33, 33, 1, 1, 1, 1,
       32, 32);
   generator.Add<TanHLayer<> >();
 
-  // Create GAN
+  // Create DCGAN
   GaussianInitialization gaussian(0, 1);
   Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
       tolerance, shuffle);
   std::function<double()> noiseFunction = [] () {
       return math::RandNormal(0, 1);};
-  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
-      std::function<double()> > gan(trainData, generator, discriminator,
+  GAN<FFN<CrossEntropyError<> >, GaussianInitialization,
+      std::function<double()>, DCGAN> dcgan(trainData, generator, discriminator,
       gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
       discriminatorPreTrain, multiplier);
 
   Log::Info << "Training..." << std::endl;
-  gan.Train(optimizer);
+  dcgan.Train(optimizer);
 
   // Generate samples
   Log::Info << "Sampling..." << std::endl;

src/mlpack/tests/gan_test.cpp

@@ -1,5 +1,5 @@
 /**
- * @file gan_network_test.cpp
+ * @file gan_test.cpp
  * @author Kris Singh
  * @author Shikhar Jaiswal
  *
@@ -14,9 +14,7 @@
 
 #include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
 #include <mlpack/methods/ann/loss_functions/cross_entropy_error.hpp>
-#include <mlpack/methods/ann/loss_functions/sigmoid_cross_entropy_error.hpp>
 #include <mlpack/methods/ann/gan.hpp>
-#include <mlpack/methods/ann/ffn.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/softmax_regression/softmax_regression.hpp>
 #include <mlpack/core/optimizers/adam/adam.hpp>
@@ -166,8 +164,7 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
   trainData.load("mnist_first250_training_4s_and_9s.arm");
   Log::Info << arma::size(trainData) << std::endl;
 
-  if (datasetMaxCols > 0)
-    trainData = trainData.cols(0, datasetMaxCols - 1);
+  trainData = trainData.cols(0, datasetMaxCols - 1);
 
   size_t numIterations = trainData.n_cols * numEpoches;
   numIterations /= batchSize;
@@ -177,7 +174,7 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
   Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;
 
   // Create the Discriminator network
-  FFN<SigmoidCrossEntropyError<> > discriminator;
+  FFN<CrossEntropyError<> > discriminator;
   discriminator.Add<Convolution<> >(1, dNumKernels, 5, 5, 1, 1, 2, 2, 28, 28);
   discriminator.Add<ReLULayer<> >();
   discriminator.Add<MeanPooling<> >(2, 2, 2, 2);
@@ -190,14 +187,17 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
   discriminator.Add<Linear<> >(1024, 1);
 
   // Create the Generator network
-  FFN<SigmoidCrossEntropyError<> > generator;
+  FFN<CrossEntropyError<> > generator;
   generator.Add<Linear<> >(noiseDim, 3136);
+  generator.Add<BatchNorm<> >(3136);
   generator.Add<ReLULayer<> >();
   generator.Add<Convolution<> >(1, noiseDim / 2, 3, 3, 2, 2, 1, 1, 56, 56);
+  generator.Add<BatchNorm<> >(39200);
   generator.Add<ReLULayer<> >();
   generator.Add<BilinearInterpolation<> >(28, 28, 56, 56, noiseDim / 2);
   generator.Add<Convolution<> >(noiseDim / 2, noiseDim / 4, 3, 3, 2, 2, 1, 1,
       56, 56);
+  generator.Add<BatchNorm<> >(19600);
   generator.Add<ReLULayer<> >();
   generator.Add<BilinearInterpolation<> >(28, 28, 56, 56, noiseDim / 4);
   generator.Add<Convolution<> >(noiseDim / 4, 1, 3, 3, 2, 2, 1, 1, 56, 56);
@@ -209,7 +209,7 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
       tolerance, shuffle);
   std::function<double()> noiseFunction = [] () {
       return math::RandNormal(0, 1);};
-  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
+  GAN<FFN<CrossEntropyError<> >, GaussianInitialization,
       std::function<double()> > gan(trainData, generator, discriminator,
       gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
       discriminatorPreTrain, multiplier);

src/mlpack/tests/loss_functions_test.cpp

@@ -10,6 +10,7 @@
  * 3-clause BSD license along with mlpack.  If not, see
  * http://www.opensource.org/licenses/BSD-3-Clause for more information.
  */
+#include <mlpack/methods/ann/loss_functions/earth_mover_distance.hpp>
 #include <mlpack/methods/ann/loss_functions/kl_divergence.hpp>
 #include <mlpack/methods/ann/loss_functions/mean_squared_error.hpp>
 #include <mlpack/methods/ann/loss_functions/sigmoid_cross_entropy_error.hpp>
@@ -169,7 +170,7 @@ BOOST_AUTO_TEST_CASE(SimpleSigmoidCrossEntropyLayerTest)
   SigmoidCrossEntropyError<> module;
 
   // Test the Forward function on a user generator input and compare it against
-  // the manually calculated result.
+  // the calculated result.
   input1 = arma::mat("0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5");
   target1 = arma::zeros(1, 8);
   double error1 = module.Forward(std::move(input1), std::move(target1));
@@ -219,4 +220,42 @@ BOOST_AUTO_TEST_CASE(SimpleSigmoidCrossEntropyLayerTest)
   BOOST_REQUIRE_EQUAL(output.n_cols, input3.n_cols);
 }
 
+/**
+ * Simple test for the Earth Mover Distance Layer.
+ */
+BOOST_AUTO_TEST_CASE(SimpleEarthMoverDistanceLayerTest)
+{
+  arma::mat input1, input2, output, target1, target2, expectedOutput;
+  EarthMoverDistance<> module;
+
+  // Test the Forward function on a user generator input and compare it against
+  // the manually calculated result.
+  input1 = arma::mat("0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5");
+  target1 = arma::zeros(1, 8);
+  double error1 = module.Forward(std::move(input1), std::move(target1));
+  double expected = 0.0;
+  BOOST_REQUIRE_SMALL(error1 / input1.n_elem - expected, 1e-7);
+
+  input2 = arma::mat("1 2 3 4 5");
+  target2 = arma::mat("1 0 1 0 1");
+  double error2 = module.Forward(std::move(input2), std::move(target2));
+  expected = -1.8;
+  BOOST_REQUIRE_SMALL(error2 / input2.n_elem - expected, 1e-6);
+
+  // Test the Backward function.
+  module.Backward(std::move(input1), std::move(target1), std::move(output));
+  expected = 0.0;
+  for (size_t i = 0; i < output.n_elem; i++)
+    BOOST_REQUIRE_SMALL(output(i) - expected, 1e-5);
+  BOOST_REQUIRE_EQUAL(output.n_rows, input1.n_rows);
+  BOOST_REQUIRE_EQUAL(output.n_cols, input1.n_cols);
+
+  expectedOutput = arma::mat("-1 0 -1 0 -1");
+  module.Backward(std::move(input2), std::move(target2), std::move(output));
+  for (size_t i = 0; i < output.n_elem; i++)
+    BOOST_REQUIRE_SMALL(output(i) - expectedOutput(i), 1e-5);
+  BOOST_REQUIRE_EQUAL(output.n_rows, input2.n_rows);
+  BOOST_REQUIRE_EQUAL(output.n_cols, input2.n_cols);
+}
+
 BOOST_AUTO_TEST_SUITE_END();

src/mlpack/tests/wgan_test.cpp

@@ -0,0 +1,288 @@
+/**
+ * @file wgan_test.cpp
+ * @author Shikhar Jaiswal
+ *
+ * Tests the WGAN network.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#include <mlpack/core.hpp>
+
+#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
+#include <mlpack/methods/ann/loss_functions/earth_mover_distance.hpp>
+#include <mlpack/methods/ann/gan.hpp>
+#include <mlpack/methods/ann/ffn.hpp>
+#include <mlpack/methods/ann/layer/layer.hpp>
+#include <mlpack/methods/softmax_regression/softmax_regression.hpp>
+#include <mlpack/core/optimizers/adam/adam.hpp>
+
+#include <boost/test/unit_test.hpp>
+#include "test_tools.hpp"
+
+using namespace mlpack;
+using namespace mlpack::ann;
+using namespace mlpack::math;
+using namespace mlpack::optimization;
+using namespace mlpack::regression;
+using namespace std::placeholders;
+
+BOOST_AUTO_TEST_SUITE(WGANNetworkTest);
+
+/*
+ * Tests the standard WGAN implementation on the MNIST dataset.
+ * It's not viable to train on bigger parameters due to time constraints.
+ * Please refer mlpack/models repository for the tutorial.
+ */
+BOOST_AUTO_TEST_CASE(WGANMNISTTest)
+{
+  size_t dNumKernels = 32;
+  size_t discriminatorPreTrain = 5;
+  size_t batchSize = 5;
+  size_t noiseDim = 100;
+  size_t generatorUpdateStep = 1;
+  size_t numSamples = 10;
+  double stepSize = 0.0003;
+  double eps = 1e-8;
+  size_t numEpoches = 1;
+  double tolerance = 1e-5;
+  int datasetMaxCols = 10;
+  bool shuffle = true;
+  double multiplier = 10;
+  double clippingParameter = 0.01;
+
+  Log::Info << std::boolalpha
+      << " batchSize = " << batchSize << std::endl
+      << " generatorUpdateStep = " << generatorUpdateStep << std::endl
+      << " noiseDim = " << noiseDim << std::endl
+      << " numSamples = " << numSamples << std::endl
+      << " stepSize = " << stepSize << std::endl
+      << " numEpoches = " << numEpoches << std::endl
+      << " tolerance = " << tolerance << std::endl
+      << " shuffle = " << shuffle << std::endl;
+
+  arma::mat trainData;
+  trainData.load("mnist_first250_training_4s_and_9s.arm");
+  Log::Info << arma::size(trainData) << std::endl;
+
+  trainData = trainData.cols(0, datasetMaxCols - 1);
+
+  size_t numIterations = trainData.n_cols * numEpoches;
+  numIterations /= batchSize;
+
+  Log::Info << "Dataset loaded (" << trainData.n_rows << ", "
+            << trainData.n_cols << ")" << std::endl;
+  Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;
+
+  // Create the Discriminator network
+  FFN<EarthMoverDistance<> > discriminator;
+  discriminator.Add<Convolution<> >(1, dNumKernels, 4, 4, 2, 2, 1, 1, 28, 28);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(dNumKernels, 2 * dNumKernels, 4, 4, 2, 2,
+      1, 1, 14, 14);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(2 * dNumKernels, 4 * dNumKernels, 4, 4,
+      2, 2, 1, 1, 7, 7);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(4 * dNumKernels, 8 * dNumKernels, 4, 4,
+      2, 2, 2, 2, 3, 3);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(8 * dNumKernels, 1, 4, 4, 1, 1,
+      1, 1, 2, 2);
+  discriminator.Add<SigmoidLayer<> >();
+
+  // Create the Generator network
+  FFN<EarthMoverDistance<> > generator;
+  generator.Add<TransposedConvolution<> >(noiseDim, 8 * dNumKernels, 2, 2,
+      1, 1, 1, 1, 1, 1);
+  generator.Add<BatchNorm<> >(1024);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(8 * dNumKernels, 4 * dNumKernels,
+      2, 2, 1, 1, 0, 0, 2, 2);
+  generator.Add<BatchNorm<> >(1152);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(4 * dNumKernels, 2 * dNumKernels,
+      5, 5, 2, 2, 1, 1, 3, 3);
+  generator.Add<BatchNorm<> >(3136);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(2 * dNumKernels, dNumKernels, 8, 8,
+      1, 1, 1, 1, 7, 7);
+  generator.Add<BatchNorm<> >(6272);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(dNumKernels, 1, 15, 15, 1, 1, 1, 1,
+      14, 14);
+  generator.Add<TanHLayer<> >();
+
+  // Create WGAN
+  GaussianInitialization gaussian(0, 1);
+  Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
+      tolerance, shuffle);
+  std::function<double()> noiseFunction = [] () {
+      return math::RandNormal(0, 1);};
+  GAN<FFN<EarthMoverDistance<> >, GaussianInitialization,
+      std::function<double()>, WGAN> wgan(trainData, generator, discriminator,
+      gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
+      discriminatorPreTrain, multiplier, clippingParameter);
+
+  Log::Info << "Training..." << std::endl;
+  wgan.Train(optimizer);
+
+  // Generate samples
+  Log::Info << "Sampling..." << std::endl;
+  arma::mat noise(noiseDim, batchSize);
+  size_t dim = std::sqrt(trainData.n_rows);
+  arma::mat generatedData(2 * dim, dim * numSamples);
+
+  for (size_t i = 0; i < numSamples; i++)
+  {
+    arma::mat samples;
+    noise.imbue( [&]() { return noiseFunction(); } );
+
+    generator.Forward(noise, samples);
+    samples.reshape(dim, dim);
+    samples = samples.t();
+
+    generatedData.submat(0, i * dim, dim - 1, i * dim + dim - 1) = samples;
+
+    samples = trainData.col(math::RandInt(0, trainData.n_cols));
+    samples.reshape(dim, dim);
+    samples = samples.t();
+
+    generatedData.submat(dim,
+        i * dim, 2 * dim - 1, i * dim + dim - 1) = samples;
+  }
+
+  Log::Info << "Output generated!" << std::endl;
+}
+
+/*
+ * Tests the gradient-penalized WGAN implementation on the MNIST dataset.
+ * It's not viable to train on bigger parameters due to time constraints.
+ * Please refer mlpack/models repository for the tutorial.
+ */
+BOOST_AUTO_TEST_CASE(WGANGPMNISTTest)
+{
+  size_t dNumKernels = 32;
+  size_t discriminatorPreTrain = 5;
+  size_t batchSize = 5;
+  size_t noiseDim = 100;
+  size_t generatorUpdateStep = 1;
+  size_t numSamples = 10;
+  double stepSize = 0.0003;
+  double eps = 1e-8;
+  size_t numEpoches = 1;
+  double tolerance = 1e-5;
+  int datasetMaxCols = 10;
+  bool shuffle = true;
+  double multiplier = 10;
+  double clippingParameter = 0.01;
+  double lambda = 10.0;
+
+  Log::Info << std::boolalpha
+      << " batchSize = " << batchSize << std::endl
+      << " generatorUpdateStep = " << generatorUpdateStep << std::endl
+      << " noiseDim = " << noiseDim << std::endl
+      << " numSamples = " << numSamples << std::endl
+      << " stepSize = " << stepSize << std::endl
+      << " numEpoches = " << numEpoches << std::endl
+      << " tolerance = " << tolerance << std::endl
+      << " shuffle = " << shuffle << std::endl;
+
+  arma::mat trainData;
+  trainData.load("mnist_first250_training_4s_and_9s.arm");
+  Log::Info << arma::size(trainData) << std::endl;
+
+  trainData = trainData.cols(0, datasetMaxCols - 1);
+
+  size_t numIterations = trainData.n_cols * numEpoches;
+  numIterations /= batchSize;
+
+  Log::Info << "Dataset loaded (" << trainData.n_rows << ", "
+            << trainData.n_cols << ")" << std::endl;
+  Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;
+
+  // Create the Discriminator network
+  FFN<EarthMoverDistance<> > discriminator;
+  discriminator.Add<Convolution<> >(1, dNumKernels, 4, 4, 2, 2, 1, 1, 28, 28);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(dNumKernels, 2 * dNumKernels, 4, 4, 2, 2,
+      1, 1, 14, 14);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(2 * dNumKernels, 4 * dNumKernels, 4, 4,
+      2, 2, 1, 1, 7, 7);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(4 * dNumKernels, 8 * dNumKernels, 4, 4,
+      2, 2, 2, 2, 3, 3);
+  discriminator.Add<LeakyReLU<> >(0.2);
+  discriminator.Add<Convolution<> >(8 * dNumKernels, 1, 4, 4, 1, 1,
+      1, 1, 2, 2);
+  discriminator.Add<SigmoidLayer<> >();
+
+  // Create the Generator network
+  FFN<EarthMoverDistance<> > generator;
+  generator.Add<TransposedConvolution<> >(noiseDim, 8 * dNumKernels, 2, 2,
+      1, 1, 1, 1, 1, 1);
+  generator.Add<BatchNorm<> >(1024);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(8 * dNumKernels, 4 * dNumKernels,
+      2, 2, 1, 1, 0, 0, 2, 2);
+  generator.Add<BatchNorm<> >(1152);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(4 * dNumKernels, 2 * dNumKernels,
+      5, 5, 2, 2, 1, 1, 3, 3);
+  generator.Add<BatchNorm<> >(3136);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(2 * dNumKernels, dNumKernels, 8, 8,
+      1, 1, 1, 1, 7, 7);
+  generator.Add<BatchNorm<> >(6272);
+  generator.Add<ReLULayer<> >();
+  generator.Add<TransposedConvolution<> >(dNumKernels, 1, 15, 15, 1, 1, 1, 1,
+      14, 14);
+  generator.Add<TanHLayer<> >();
+
+  // Create WGANGP
+  GaussianInitialization gaussian(0, 1);
+  Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
+      tolerance, shuffle);
+  std::function<double()> noiseFunction = [] () {
+      return math::RandNormal(0, 1);};
+  GAN<FFN<EarthMoverDistance<> >, GaussianInitialization,
+      std::function<double()>, WGANGP > wgan(trainData, generator,
+      discriminator, gaussian, noiseFunction, noiseDim, batchSize,
+      generatorUpdateStep, discriminatorPreTrain, multiplier, clippingParameter,
+      lambda);
+
+  Log::Info << "Training..." << std::endl;
+  wgan.Train(optimizer);
+
+  // Generate samples
+  Log::Info << "Sampling..." << std::endl;
+  arma::mat noise(noiseDim, batchSize);
+  size_t dim = std::sqrt(trainData.n_rows);
+  arma::mat generatedData(2 * dim, dim * numSamples);
+
+  for (size_t i = 0; i < numSamples; i++)
+  {
+    arma::mat samples;
+    noise.imbue( [&]() { return noiseFunction(); } );
+
+    generator.Forward(noise, samples);
+    samples.reshape(dim, dim);
+    samples = samples.t();
+
+    generatedData.submat(0, i * dim, dim - 1, i * dim + dim - 1) = samples;
+
+    samples = trainData.col(math::RandInt(0, trainData.n_cols));
+    samples.reshape(dim, dim);
+    samples = samples.t();
+
+    generatedData.submat(dim,
+        i * dim, 2 * dim - 1, i * dim + dim - 1) = samples;
+  }
+
+  Log::Info << "Output generated!" << std::endl;
+}
+
+BOOST_AUTO_TEST_SUITE_END();