Add batch support.

src/mlpack/core/optimizers/iqn/iqn.hpp

@@ -74,11 +74,13 @@ class IQN
    * equal one pass over the dataset).
    *
    * @param stepSize Step size for each iteration.
+   * @param batchSize Size of each batch.
    * @param maxIterations Maximum number of iterations allowed (0 means no
    *     limit).
    * @param tolerance Maximum absolute tolerance to terminate algorithm.
    */
   IQN(const double stepSize = 0.01,
+      const size_t batchSize = 10,
       const size_t maxIterations = 100000,
       const double tolerance = 1e-5);
 
@@ -100,6 +102,11 @@ class IQN
   //! Modify the step size.
   double& StepSize() { return stepSize; }
 
+  //! Get the batch size.
+  size_t BatchSize() const { return batchSize; }
+  //! Modify the batch size.
+  size_t& BatchSize() { return batchSize; }
+
   //! Get the maximum number of iterations (0 indicates no limit).
   size_t MaxIterations() const { return maxIterations; }
   //! Modify the maximum number of iterations (0 indicates no limit).
@@ -114,6 +121,9 @@ class IQN
   //! The step size for each example.
   double stepSize;
 
+  //! The size of each batch.
+  size_t batchSize;
+
   //! The maximum number of allowed iterations.
   size_t maxIterations;
 

src/mlpack/core/optimizers/iqn/iqn_impl.hpp

@@ -21,9 +21,11 @@ namespace mlpack {
 namespace optimization {
 
 IQN::IQN(const double stepSize,
+         const size_t batchSize,
          const size_t maxIterations,
          const double tolerance) :
     stepSize(stepSize),
+    batchSize(batchSize),
     maxIterations(maxIterations),
     tolerance(tolerance)
 { /* Nothing to do. */ }
@@ -32,27 +34,36 @@ IQN::IQN(const double stepSize,
 template<typename DecomposableFunctionType>
 double IQN::Optimize(DecomposableFunctionType& function, arma::mat& iterate)
 {
-  // Find the number of functions to use.
+  // Find the number of functions.
   const size_t numFunctions = function.NumFunctions();
+  size_t numBatches = numFunctions / batchSize;
+  if (numFunctions % batchSize != 0)
+    ++numBatches; // Capture last few.
 
   // To keep track of where we are and how things are going.
   double overallObjective = 0;
 
-  arma::cube y(iterate.n_rows, iterate.n_cols, numFunctions);
-  arma::cube t(iterate.n_elem, 1, numFunctions);
-  arma::cube Q(iterate.n_elem, iterate.n_elem, numFunctions);
+  arma::cube y(iterate.n_rows, iterate.n_cols, numBatches);
+  arma::cube t(iterate.n_elem, 1, numBatches);
+  arma::cube Q(iterate.n_elem, iterate.n_elem, numBatches);
   arma::mat initialIterate = arma::randn(iterate.n_rows, iterate.n_cols);
   arma::mat B = arma::eye(iterate.n_elem, iterate.n_elem);
 
   arma::mat g = arma::zeros(iterate.n_rows, iterate.n_cols);
-  for (size_t i = 0; i < numFunctions; ++i)
+  for (size_t i = 0, f = 0; i < numFunctions; f++)
   {
-    t.slice(i) = arma::mat(initialIterate.memptr(), iterate.n_elem,
-        1, false, false);
-    function.Gradient(initialIterate, i, y.slice(i));
-    Q.slice(i).eye();
+    // Find the effective batch size (the last batch may be smaller).
+    const size_t effectiveBatchSize = std::min(batchSize, numFunctions - i);
 
-    g += y.slice(i);
+    t.slice(f) = arma::mat(initialIterate.memptr(), iterate.n_elem,
+        1, false, false);
+    function.Gradient(initialIterate, i, y.slice(f), effectiveBatchSize);
+
+    Q.slice(f).eye();
+    g += y.slice(f);
+    y.slice(f) /= (double) effectiveBatchSize;
+
+    i += effectiveBatchSize;
   }
   g /= numFunctions;
 
@@ -66,44 +77,56 @@ double IQN::Optimize(DecomposableFunctionType& function, arma::mat& iterate)
 
   for (size_t i = 1; i != maxIterations; ++i)
   {
-    for (size_t j = 0; j < numFunctions; ++j)
+    for (size_t j = 0, f = 0; f < numFunctions; j++)
     {
-      // Cyclicly iterating through the nnumber of functions.
-      const size_t it = ((j + 1) % numFunctions);
+      // Cyclicly iterating through the number of functions.
+      const size_t it = ((j + 1) % numBatches);
+
+      // Find the effective batch size (the last batch may be smaller).
+      const size_t effectiveBatchSize = std::min(batchSize, numFunctions -
+          it * batchSize);
 
       if (arma::norm(iterateVec - t.slice(it)) > 0)
       {
-        function.Gradient(iterate, it, gradient);
+        function.Gradient(iterate, it * batchSize, gradient,
+            effectiveBatchSize);
+        gradient /= effectiveBatchSize;
+
         const arma::mat s = iterateVec - t.slice(it);
         const arma::mat yy = arma::vectorise(gradient - y.slice(it));
 
-         const arma::mat stochasticHessian = Q.slice(it) + yy * yy.t() /
-              arma::as_scalar(yy.t() * s) - Q.slice(it) * s * s.t() *
-              Q.slice(it) / arma::as_scalar(s.t() * Q.slice(it) * s);
+        const arma::mat stochasticHessian = Q.slice(it) + yy * yy.t() /
+            arma::as_scalar(yy.t() * s) - Q.slice(it) * s * s.t() *
+            Q.slice(it) / arma::as_scalar(s.t() * Q.slice(it) * s);
 
         // Update aggregate Hessian approximation.
-        B += (1.0 / numFunctions) * (stochasticHessian - Q.slice(it));
+        B += (1.0 / numBatches) * (stochasticHessian - Q.slice(it));
 
         // Update aggregate Hessian-variable product.
-        u += (1.0 / numFunctions) * (stochasticHessian * iterateVec -
+        u += (1.0 / numBatches) * (stochasticHessian * iterateVec -
             Q.slice(it) * t.slice(it));
 
         // Update aggregate gradient.
-        g += (1.0 / numFunctions) * (gradient - y.slice(it));
+        g += (1.0 / numBatches) * (gradient - y.slice(it));
 
         // Update the function information tables.
         Q.slice(it) = stochasticHessian;
         y.slice(it) = gradient;
         t.slice(it) = iterateVec;
 
-        iterateVec = stepSize * B.i() * (u - gVec) +
-            (1 - stepSize) * iterateVec;
+        iterateVec = stepSize * B.i() * (u - gVec) + (1 - stepSize) *
+            iterateVec;
       }
+
+      f+= effectiveBatchSize;
     }
 
     overallObjective = 0;
-    for (size_t i = 0; i < numFunctions; ++i)
-      overallObjective += function.Evaluate(iterate, i);
+    for (size_t f = 0; f < numFunctions; f += batchSize)
+    {
+      const size_t effectiveBatchSize = std::min(batchSize, numFunctions - f);
+      overallObjective += function.Evaluate(iterate, f, effectiveBatchSize);
+    }
     overallObjective /= numFunctions;
 
     // Output current objective function.

src/mlpack/tests/iqn_test.cpp

@@ -70,15 +70,19 @@ BOOST_AUTO_TEST_CASE(LogisticRegressionTest)
     testResponses[i] = 1;
   }
 
-  IQN iqn(0.01, 5000, 1e-3);
-  LogisticRegression<> lr(shuffledData, shuffledResponses, iqn, 0.5);
+  // Now run SGDR with snapshot ensembles on a couple of batch sizes.
+  for (size_t batchSize = 1; batchSize < 9; batchSize += 4)
+  {
+    IQN iqn(0.01, batchSize, 5000, 1e-3);
+    LogisticRegression<> lr(shuffledData, shuffledResponses, iqn, 0.5);
 
-  // Ensure that the error is close to zero.
-  const double acc = lr.ComputeAccuracy(data, responses);
-  BOOST_REQUIRE_CLOSE(acc, 100.0, 0.3); // 0.3% error tolerance.
+    // Ensure that the error is close to zero.
+    const double acc = lr.ComputeAccuracy(data, responses);
+    BOOST_REQUIRE_CLOSE(acc, 100.0, 1.3); // 1.3% error tolerance.
 
-  const double testAcc = lr.ComputeAccuracy(testData, testResponses);
-  BOOST_REQUIRE_CLOSE(testAcc, 100.0, 0.6); // 0.6% error tolerance.
+    const double testAcc = lr.ComputeAccuracy(testData, testResponses);
+    BOOST_REQUIRE_CLOSE(testAcc, 100.0, 1.6); // 1.6% error tolerance.
+  }
 }
 
 BOOST_AUTO_TEST_SUITE_END();