Add EvaluateWithGradient Function

src/mlpack/methods/ann/CMakeLists.txt

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

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

@@ -0,0 +1,18 @@
+# Define the files we need to compile
+# Anything not in this list will not be compiled into mlpack.
+set(SOURCES
+  gan.hpp
+  gan_impl.hpp
+  gan_policies.hpp
+  wgan_impl.hpp
+  wgangp_impl.hpp
+)
+
+# Add directory name to sources.
+set(DIR_SRCS)
+foreach(file ${SOURCES})
+    set(DIR_SRCS ${DIR_SRCS} ${CMAKE_CURRENT_SOURCE_DIR}/${file})
+endforeach()
+# Append sources (with directory name) to list of all mlpack sources (used at
+# the parent scope).
+set(MLPACK_SRCS ${MLPACK_SRCS} ${DIR_SRCS} PARENT_SCOPE)

src/mlpack/methods/ann/gan/gan.hpp

@@ -8,13 +8,13 @@
  * 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_HPP
-#define MLPACK_METHODS_ANN_GAN_HPP
+#ifndef MLPACK_METHODS_ANN_GAN_GAN_HPP
+#define MLPACK_METHODS_ANN_GAN_GAN_HPP
 
 #include <mlpack/core.hpp>
 
 #include <mlpack/methods/ann/ffn.hpp>
-#include <mlpack/methods/ann/gan_policies.hpp>
+#include <mlpack/methods/ann/gan/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,6 +31,8 @@ namespace ann /** Artificial Neural Network. **/ {
  * networks contesting with each other in a zero-sum game framework. This
  * technique can generate photographs that look at least superficially
  * authentic to human observers, having many realistic characteristics.
+ * GANs have been used in Text-to-Image Synthesis, Medical Drug Discovery,
+ * High Resolution Imagery Generation, Neural Machine Translation and so on.
  *
  * For more information, see the following paper:
  *
@@ -75,17 +77,23 @@ class GAN
    * @param lambda Parameter for setting the gradient penalty.
    */
   GAN(arma::mat& trainData,
-      Model& generator,
-      Model& discriminator,
-      InitializationRuleType initializeRule,
-      Noise noiseFunction,
-      size_t noiseDim,
-      size_t batchSize,
-      size_t generatorUpdateStep,
-      size_t preTrainSize,
-      double multiplier,
-      double clippingParameter = 0.01,
-      double lambda = 10.0);
+      Model generator,
+      Model discriminator,
+      InitializationRuleType& initializeRule,
+      Noise& noiseFunction,
+      const size_t noiseDim,
+      const size_t batchSize,
+      const size_t generatorUpdateStep,
+      const size_t preTrainSize,
+      const double multiplier,
+      const double clippingParameter = 0.01,
+      const double lambda = 10.0);
+
+  //! Copy constructor.
+  GAN(const GAN&);
+
+  //! Move constructor.
+  GAN(GAN&&);
 
   // Reset function.
   void Reset();
@@ -140,6 +148,60 @@ class GAN
            const size_t i,
            const size_t batchSize);
 
+  /**
+   * EvaluateWithGradient function for the Standard GAN and DCGAN.
+   * This function gives the performance of the Standard GAN or DCGAN on the
+   * current input, while updating Gradients.
+   *
+   * @param parameters The parameters of the network.
+   * @param i Index of the current input.
+   * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename GradType, typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, StandardGAN>::value ||
+                          std::is_same<Policy, DCGAN>::value, double>::type
+  EvaluateWithGradient(const arma::mat& parameters,
+                       const size_t i,
+                       GradType& gradient,
+                       const size_t batchSize);
+
+  /**
+   * EvaluateWithGradient function for the WGAN.
+   * This function gives the performance of the WGAN on the
+   * current input, while updating Gradients.
+   *
+   * @param parameters The parameters of the network.
+   * @param i Index of the current input.
+   * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename GradType, typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGAN>::value,
+                          double>::type
+  EvaluateWithGradient(const arma::mat& parameters,
+                       const size_t i,
+                       GradType& gradient,
+                       const size_t batchSize);
+
+  /**
+   * EvaluateWithGradient function for the WGAN-GP.
+   * This function gives the performance of the WGAN-GP on the
+   * current input, while updating Gradients.
+   *
+   * @param parameters The parameters of the network.
+   * @param i Index of the current input.
+   * @param gradient Variable to store the present gradient.
+   * @param batchSize Variable to store the present number of inputs.
+   */
+  template<typename GradType, typename Policy = PolicyType>
+  typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                          double>::type
+  EvaluateWithGradient(const arma::mat& parameters,
+                       const size_t i,
+                       GradType& gradient,
+                       const size_t batchSize);
+
   /**
    * Gradient function for Standard GAN and DCGAN.
    * This function passes the gradient based on which network is being
@@ -217,9 +279,18 @@ class GAN
 
   //! Return the parameters of the network.
   const arma::mat& Parameters() const { return parameter; }
-  //! Modify the parameters of the network
+  //! Modify the parameters of the network.
   arma::mat& Parameters() { return parameter; }
 
+  //! Return the generator of the GAN.
+  const Model& Generator() const { return generator; }
+  //! Modify the generator of the GAN.
+  Model& Generator() { return generator; }
+  //! Return the discriminator of the GAN.
+  const Model& Discriminator() const { return discriminator; }
+  //! Modify the discriminator of the GAN.
+  Model& Discriminator() { return discriminator; }
+
   //! Return the number of separable functions (the number of predictor points).
   size_t NumFunctions() const { return numFunctions; }
 
@@ -233,9 +304,9 @@ class GAN
   //! Locally stored parameters of the network.
   arma::mat parameter;
   //! Locally stored Generator network.
-  Model& generator;
+  Model generator;
   //! Locally stored Discriminator network.
-  Model& discriminator;
+  Model discriminator;
   //! Locally stored Initializer.
   InitializationRuleType initializeRule;
   //! Locally stored Noise function
@@ -296,5 +367,8 @@ class GAN
 
 // Include implementation.
 #include "gan_impl.hpp"
+#include "wgan_impl.hpp"
+#include "wgangp_impl.hpp"
+
 
 #endif

src/mlpack/methods/ann/gan/gan_impl.hpp

@@ -8,8 +8,8 @@
  * 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_IMPL_HPP
-#define MLPACK_METHODS_ANN_GAN_IMPL_HPP
+#ifndef MLPACK_METHODS_ANN_GAN_GAN_IMPL_HPP
+#define MLPACK_METHODS_ANN_GAN_GAN_IMPL_HPP
 
 #include "gan.hpp"
 
@@ -31,20 +31,20 @@ template<
 >
 GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
     arma::mat& predictors,
-    Model& generator,
-    Model& discriminator,
-    InitializationRuleType initializeRule,
-    Noise noiseFunction,
-    size_t noiseDim,
-    size_t batchSize,
-    size_t generatorUpdateStep,
-    size_t preTrainSize,
-    double multiplier,
-    double clippingParameter,
-    double lambda):
+    Model generator,
+    Model discriminator,
+    InitializationRuleType& initializeRule,
+    Noise& noiseFunction,
+    const size_t noiseDim,
+    const size_t batchSize,
+    const size_t generatorUpdateStep,
+    const size_t preTrainSize,
+    const double multiplier,
+    const double clippingParameter,
+    const double lambda):
     predictors(predictors),
-    generator(generator),
-    discriminator(discriminator),
+    generator(std::move(generator)),
+    discriminator(std::move(discriminator)),
     initializeRule(initializeRule),
     noiseFunction(noiseFunction),
     noiseDim(noiseDim),
@@ -57,33 +57,95 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
     reset(false)
 {
   // Insert IdentityLayer for joining the Generator and Discriminator.
-  discriminator.network.insert(
-      discriminator.network.begin(),
+  this->discriminator.network.insert(
+      this->discriminator.network.begin(),
       new IdentityLayer<>());
 
   counter = 0;
   currentBatch = 0;
 
-  discriminator.deterministic = generator.deterministic = true;
+  this->discriminator.deterministic = this->generator.deterministic = true;
 
   responses.set_size(1, predictors.n_cols);
   responses.ones();
 
-  discriminator.predictors.set_size(predictors.n_rows,
+  this->discriminator.predictors.set_size(predictors.n_rows,
       predictors.n_cols + batchSize);
-  discriminator.predictors.cols(0, predictors.n_cols - 1) = predictors;
+  this->discriminator.predictors.cols(0, predictors.n_cols - 1) = predictors;
 
-  discriminator.responses.set_size(1, predictors.n_cols + batchSize);
-  discriminator.responses.ones();
-  discriminator.responses.cols(predictors.n_cols,
+  this->discriminator.responses.set_size(1, predictors.n_cols + batchSize);
+  this->discriminator.responses.ones();
+  this->discriminator.responses.cols(predictors.n_cols,
       predictors.n_cols + batchSize - 1) = arma::zeros(1, batchSize);
 
   numFunctions = predictors.n_cols;
 
   noise.set_size(noiseDim, batchSize);
 
-  generator.predictors.set_size(noiseDim, batchSize);
-  generator.responses.set_size(predictors.n_rows, batchSize);
+  this->generator.predictors.set_size(noiseDim, batchSize);
+  this->generator.responses.set_size(predictors.n_rows, batchSize);
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
+    const GAN& network):
+    predictors(network.predictors),
+    responses(network.responses),
+    generator(network.generator),
+    discriminator(network.discriminator),
+    initializeRule(network.initializeRule),
+    noiseFunction(network.noiseFunction),
+    noiseDim(network.noiseDim),
+    batchSize(network.batchSize),
+    generatorUpdateStep(network.generatorUpdateStep),
+    preTrainSize(network.preTrainSize),
+    multiplier(network.multiplier),
+    clippingParameter(network.clippingParameter),
+    lambda(network.lambda),
+    reset(network.reset),
+    counter(network.counter),
+    currentBatch(network.currentBatch),
+    parameter(network.parameter),
+    numFunctions(network.numFunctions),
+    noise(network.noise)
+{
+  /* Nothing to do here */
+}
+
+template<
+  typename Model,
+  typename InitializationRuleType,
+  typename Noise,
+  typename PolicyType
+>
+GAN<Model, InitializationRuleType, Noise, PolicyType>::GAN(
+    GAN&& network):
+    predictors(std::move(network.predictors)),
+    responses(std::move(network.responses)),
+    generator(std::move(network.generator)),
+    discriminator(std::move(network.discriminator)),
+    initializeRule(std::move(network.initializeRule)),
+    noiseFunction(std::move(network.noiseFunction)),
+    noiseDim(network.noiseDim),
+    batchSize(network.batchSize),
+    generatorUpdateStep(network.generatorUpdateStep),
+    preTrainSize(network.preTrainSize),
+    multiplier(network.multiplier),
+    clippingParameter(network.clippingParameter),
+    lambda(network.lambda),
+    reset(network.reset),
+    counter(network.counter),
+    currentBatch(network.currentBatch),
+    parameter(std::move(network.parameter)),
+    numFunctions(network.numFunctions),
+    noise(std::move(network.noise))
+{
+  /* Nothing to do here */
 }
 
 template<
@@ -193,16 +255,44 @@ template<
   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 */)
+template<typename GradType, 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>::
+EvaluateWithGradient(const arma::mat& /* parameters */,
+                     const size_t i,
+                     GradType& 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);
+
   currentInput = arma::mat(predictors.memptr() + (i * predictors.n_rows),
       predictors.n_rows, batchSize, false, false);
   currentTarget = arma::mat(responses.memptr() + i, 1, batchSize, false,
@@ -214,16 +304,18 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
       outputParameterVisitor,
       discriminator.network.back())), std::move(currentTarget));
 
+  // 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.Forward(std::move(discriminator.predictors.cols(numFunctions,
       numFunctions + batchSize - 1)));
   discriminator.responses.cols(numFunctions, numFunctions + batchSize - 1) =
-      -arma::ones(1, batchSize);
-
+      arma::zeros(1, batchSize);
   currentTarget = arma::mat(discriminator.responses.memptr() + numFunctions,
       1, batchSize, false, false);
   res += discriminator.outputLayer.Forward(
@@ -231,65 +323,41 @@ GAN<Model, InitializationRuleType, Noise, PolicyType>::Evaluate(
       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);
+      noiseGradientDiscriminator, batchSize);
+  gradientDiscriminator += noiseGradientDiscriminator;
+
+  if (currentBatch % generatorUpdateStep == 0 && preTrainSize == 0)
+  {
+    // Minimize -log(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--;
+  }
 
   return res;
 }
@@ -385,189 +453,6 @@ Gradient(const arma::mat& /* parameters */,
   }
 }
 
-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,

src/mlpack/methods/ann/gan/gan_policies.hpp

@@ -9,8 +9,8 @@
  * 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
+#ifndef MLPACK_METHODS_ANN_GAN_GAN_POLICIES_HPP
+#define MLPACK_METHODS_ANN_GAN_GAN_POLICIES_HPP
 
 namespace mlpack {
 namespace ann /** Artificial Neural Network. */ {

src/mlpack/methods/ann/gan/wgan_impl.hpp

@@ -0,0 +1,279 @@
+/**
+ * @file wgan_impl.hpp
+ * @author Shikhar Jaiswal
+ *
+ * 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_WGAN_IMPL_HPP
+#define MLPACK_METHODS_ANN_GAN_WGAN_IMPL_HPP
+
+#include "gan.hpp"
+
+#include <mlpack/core.hpp>
+
+#include <mlpack/methods/ann/ffn.hpp>
+#include <mlpack/methods/ann/init_rules/network_init.hpp>
+#include <mlpack/methods/ann/visitor/output_parameter_visitor.hpp>
+#include <mlpack/methods/ann/activation_functions/softplus_function.hpp>
+
+namespace mlpack {
+namespace ann /** Artifical Neural Network.  */ {
+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 GradType, typename Policy>
+typename std::enable_if<std::is_same<Policy, WGAN>::value, double>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::
+EvaluateWithGradient(const arma::mat& /* parameters */,
+                     const size_t i,
+                     GradType& 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);
+
+  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));
+
+  // 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.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));
+
+  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--;
+  }
+
+  return res;
+}
+
+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--;
+  }
+}
+
+} // namespace ann
+} // namespace mlpack
+# endif

src/mlpack/methods/ann/gan/wgangp_impl.hpp

@@ -0,0 +1,302 @@
+/**
+ * @file wgangp_impl.hpp
+ * @author Shikhar Jaiswal
+ *
+ * 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_WGANGP_IMPL_HPP
+#define MLPACK_METHODS_ANN_GAN_WGANGP_IMPL_HPP
+
+#include "gan.hpp"
+
+#include <mlpack/core.hpp>
+
+#include <mlpack/methods/ann/ffn.hpp>
+#include <mlpack/methods/ann/init_rules/network_init.hpp>
+#include <mlpack/methods/ann/visitor/output_parameter_visitor.hpp>
+#include <mlpack/methods/ann/activation_functions/softplus_function.hpp>
+
+namespace mlpack {
+namespace ann /** Artifical Neural Network.  */ {
+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 GradType, typename Policy>
+typename std::enable_if<std::is_same<Policy, WGANGP>::value,
+                        double>::type
+GAN<Model, InitializationRuleType, Noise, PolicyType>::
+EvaluateWithGradient(const arma::mat& /* parameters */,
+                     const size_t i,
+                     GradType& 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);
+
+  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));
+
+  // Get the gradients of the Discriminator.
+  discriminator.Gradient(discriminator.parameter, i, gradientDiscriminator,
+      batchSize);
+
+  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);
+
+  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--;
+  }
+
+  return res;
+}
+
+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--;
+  }
+}
+
+} // namespace ann
+} // namespace mlpack
+# endif

src/mlpack/tests/dcgan_test.cpp

@@ -13,7 +13,7 @@
 
 #include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
 #include <mlpack/methods/ann/loss_functions/cross_entropy_error.hpp>
-#include <mlpack/methods/ann/gan.hpp>
+#include <mlpack/methods/ann/gan/gan.hpp>
 #include <mlpack/methods/ann/ffn.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/softmax_regression/softmax_regression.hpp>
@@ -139,7 +139,7 @@ BOOST_AUTO_TEST_CASE(DCGANMNISTTest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    dcgan.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();
 
@@ -265,7 +265,7 @@ BOOST_AUTO_TEST_CASE(DCGANCelebATest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    dcgan.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();
 

src/mlpack/tests/gan_test.cpp

@@ -14,7 +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/gan.hpp>
+#include <mlpack/methods/ann/gan/gan.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/softmax_regression/softmax_regression.hpp>
 #include <mlpack/core/optimizers/adam/adam.hpp>
@@ -92,7 +92,7 @@ BOOST_AUTO_TEST_CASE(GANTest)
 
   // Generate samples
   Log::Info << "Sampling..." << std::endl;
-  arma::mat noise(noiseDim, 1);
+  arma::mat noise(noiseDim, batchSize);
 
   size_t dim = std::sqrt(trainData.n_rows);
   arma::mat generatedData(2 * dim, dim * numSamples);
@@ -102,7 +102,7 @@ BOOST_AUTO_TEST_CASE(GANTest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    gan.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();
 
@@ -228,7 +228,7 @@ BOOST_AUTO_TEST_CASE(GANMNISTTest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    gan.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();
 

src/mlpack/tests/wgan_test.cpp

@@ -13,7 +13,7 @@
 
 #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/gan/gan.hpp>
 #include <mlpack/methods/ann/ffn.hpp>
 #include <mlpack/methods/ann/layer/layer.hpp>
 #include <mlpack/methods/softmax_regression/softmax_regression.hpp>
@@ -140,7 +140,7 @@ BOOST_AUTO_TEST_CASE(WGANMNISTTest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    wgan.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();
 
@@ -249,13 +249,13 @@ BOOST_AUTO_TEST_CASE(WGANGPMNISTTest)
   std::function<double()> noiseFunction = [] () {
       return math::RandNormal(0, 1);};
   GAN<FFN<EarthMoverDistance<> >, GaussianInitialization,
-      std::function<double()>, WGANGP > wgan(trainData, generator,
+      std::function<double()>, WGANGP > wganGP(trainData, generator,
       discriminator, gaussian, noiseFunction, noiseDim, batchSize,
       generatorUpdateStep, discriminatorPreTrain, multiplier, clippingParameter,
       lambda);
 
   Log::Info << "Training..." << std::endl;
-  wgan.Train(optimizer);
+  wganGP.Train(optimizer);
 
   // Generate samples
   Log::Info << "Sampling..." << std::endl;
@@ -268,7 +268,7 @@ BOOST_AUTO_TEST_CASE(WGANGPMNISTTest)
     arma::mat samples;
     noise.imbue( [&]() { return noiseFunction(); } );
 
-    generator.Forward(noise, samples);
+    wganGP.Generator().Forward(noise, samples);
     samples.reshape(dim, dim);
     samples = samples.t();