changes acc to suggestion

src/mlpack/methods/ann/layer/batch_norm.hpp

@@ -54,13 +54,28 @@ template <typename MatType = arma::mat>
 class BatchNormType : public Layer<MatType>
 {
  public:
-  //! Create the BatchNorm object. This will initialize maxAxis to 2.
-  //! (or effectively 1 or 0 if there are fewer input Dimensions than 3.)
+  /**
+   * Create the BatchNorm object. 
+   * 
+   * This sets the minimum and maximum axis for batch normalization to 2; so, 
+   * e.g., we apply batch normalization to the first 3 dimensions with number 
+   * of channels equal to 3rd dimension, and higher dimensions are left 
+   * untouched.  As an example, if we have a 3-dimensional input (call the 
+   * three dimensions rows, columns and slices), and `minAxis` & `maxAxis` is 
+   * 2, then we apply the same normalization across different slices.
+   */
   BatchNormType();
 
   /**
-   * Create the BatchNorm layer object for a specified number of input units.
+   * Create the BatchNorm layer object for a specified axis of input units as 
+   * channels.
+   * 
+   * As an example, if we have a 3-dimensional input (call the three dimensions
+   * rows, columns and slices), and `minAxis` is 1 & `maxAxis` is 2, then the 
+   * number of channels is equal to `columns * slices`.
    *
+   * @param minAxis The min axis along which BatchNorm is applied. Before that,
+   *                it will be treated as input point.
    * @param maxAxis The max axis along which BatchNorm is applied. After that,
    *                it will be treated as another higher dimension point.
    * @param eps The epsilon added to variance to ensure numerical stability.
@@ -68,7 +83,8 @@ class BatchNormType : public Layer<MatType>
    *                updating the parameters or momentum is used.
    * @param momentum Parameter used to to update the running mean and variance.
    */
-  BatchNormType(const size_t maxAxis,
+  BatchNormType(const size_t minAxis,
+                const size_t maxAxis,
                 const double eps = 1e-8,
                 const bool average = true,
                 const double momentum = 0.1);
@@ -179,6 +195,9 @@ class BatchNormType : public Layer<MatType>
   void serialize(Archive& ar, const uint32_t /* version */);
 
  private:
+  //! Locally-stored minAxis along which BatchNorm will apply.
+  size_t minAxis;
+
   //! Locally-stored maxAxis along which BatchNorm will apply.
   size_t maxAxis;
 

src/mlpack/methods/ann/layer/batch_norm_impl.hpp

@@ -25,6 +25,7 @@ namespace ann /** Artificial Neural Network. */ {
 template<typename MatType>
 BatchNormType<MatType>::BatchNormType() : 
     Layer<MatType>(),
+    minAxis(2),
     maxAxis(2),
     eps(1e-8),
     average(true),
@@ -39,11 +40,13 @@ BatchNormType<MatType>::BatchNormType() :
 
 template <typename MatType>
 BatchNormType<MatType>::BatchNormType(
+    const size_t minAxis,
     const size_t maxAxis,
     const double eps,
     const bool average,
     const double momentum) : 
     Layer<MatType>(),
+    minAxis(minAxis),
     maxAxis(maxAxis),
     eps(eps),
     average(average),
@@ -60,6 +63,7 @@ BatchNormType<MatType>::BatchNormType(
 template<typename MatType>
 BatchNormType<MatType>::BatchNormType(const BatchNormType& layer) :
     Layer<MatType>(layer),
+    minAxis(layer.minAxis),
     maxAxis(layer.maxAxis),
     eps(layer.eps),
     average(layer.average),
@@ -79,6 +83,7 @@ BatchNormType<MatType>::BatchNormType(const BatchNormType& layer) :
 template<typename MatType>
 BatchNormType<MatType>::BatchNormType(BatchNormType&& layer) :
     Layer<MatType>(std::move(layer)),
+    minAxis(std::move(layer.minAxis)),
     maxAxis(std::move(layer.maxAxis)),
     eps(std::move(layer.eps)),
     average(std::move(layer.average)),
@@ -101,6 +106,7 @@ BatchNormType<MatType>::operator=(const BatchNormType& layer)
   if (&layer != this)
   {
     Layer<MatType>::operator=(layer);
+    minAxis = layer.minAxis;
     maxAxis = layer.maxAxis;
     eps = layer.eps;
     average = layer.average;
@@ -125,6 +131,7 @@ BatchNormType<MatType>::operator=(
   if (&layer != this)
   {
     Layer<MatType>::operator=(std::move(layer));
+    minAxis = std::move(layer.minAxis);
     maxAxis = std::move(layer.maxAxis);
     eps = std::move(layer.eps);
     average = std::move(layer.average);
@@ -332,14 +339,22 @@ void BatchNormType<MatType>::Gradient(
 template<typename MatType>
 void BatchNormType<MatType>::ComputeOutputDimensions()
 {
+  if (minAxis > maxAxis)
+  {
+    Log::Fatal << "BatchNorm: minAxis must be less than or equal to maxAxis."
+        << std::endl;
+  }
   this->outputDimensions = this->inputDimensions;
-  size_t mainAxis = std::min(this->inputDimensions.size() - 1, maxAxis);
+  size_t mainMinAxis = std::min(this->inputDimensions.size() - 1, minAxis);
+  size_t mainMaxAxis = std::min(this->inputDimensions.size() - 1, maxAxis);
   inputDimension = 1;
-  for (size_t i = 0; i < mainAxis; i++)
+  for (size_t i = 0; i < mainMinAxis; i++)
     inputDimension *= this->inputDimensions[i];
-  size = this->inputDimensions[mainAxis];
+  size = this->inputDimensions[mainMinAxis];
+  for (size_t i = mainMinAxis + 1; i <= mainMaxAxis; i++)
+    size *= this->inputDimensions[i];
   higherDimension = 1;
-  for (size_t i = mainAxis + 1; i < this->inputDimensions.size(); i++)
+  for (size_t i = mainMaxAxis + 1; i < this->inputDimensions.size(); i++)
     higherDimension *= this->inputDimensions[i];
 }
 
@@ -350,6 +365,7 @@ void BatchNormType<MatType>::serialize(
 {
   ar(cereal::base_class<Layer<MatType>>(this));
 
+  ar(CEREAL_NVP(minAxis));
   ar(CEREAL_NVP(maxAxis));
   ar(CEREAL_NVP(eps));
   ar(CEREAL_NVP(count));

src/mlpack/tests/ann_layer_test.cpp

@@ -3029,7 +3029,7 @@ TEST_CASE("BatchNormTest", "[ANNLayerTest]")
   module1.SetWeights((double*) moduleParams.memptr());
 
   // BatchNorm layer with average parameter set to false (using momentum).
-  BatchNorm module2(2, 1e-5, false);
+  BatchNorm module2(2, 2, 1e-5, false);
   module2.Training() = true;
   module2.InputDimensions() = std::vector<size_t>({ 3, 3 });
   module2.ComputeOutputDimensions();
@@ -5471,7 +5471,7 @@ TEST_CASE("BatchNormWithMinBatchesTest", "[ANNLayerTest]")
              { 0.2507 , -0.8486 , 0.8293 } };
 
   // Check correctness of batch normalization.
-  BatchNorm module1(2, 1e-5, false, 0.1);
+  BatchNorm module1(2, 2, 1e-5, false, 0.1);
   module1.Training() = true;
   module1.InputDimensions() = std::vector<size_t>({ 1, 4, 2 });
   module1.ComputeOutputDimensions();