diff --git a/src/mlpack/methods/ann/layer/batch_norm.hpp b/src/mlpack/methods/ann/layer/batch_norm.hpp index 21594435d9..4289a435a1 100644 --- a/src/mlpack/methods/ann/layer/batch_norm.hpp +++ b/src/mlpack/methods/ann/layer/batch_norm.hpp @@ -54,13 +54,28 @@ template class BatchNormType : public Layer { 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 * 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 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; diff --git a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp index 007fee8578..af6453db74 100644 --- a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp +++ b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp @@ -25,6 +25,7 @@ namespace ann /** Artificial Neural Network. */ { template BatchNormType::BatchNormType() : Layer(), + minAxis(2), maxAxis(2), eps(1e-8), average(true), @@ -39,11 +40,13 @@ BatchNormType::BatchNormType() : template BatchNormType::BatchNormType( + const size_t minAxis, const size_t maxAxis, const double eps, const bool average, const double momentum) : Layer(), + minAxis(minAxis), maxAxis(maxAxis), eps(eps), average(average), @@ -60,6 +63,7 @@ BatchNormType::BatchNormType( template BatchNormType::BatchNormType(const BatchNormType& layer) : Layer(layer), + minAxis(layer.minAxis), maxAxis(layer.maxAxis), eps(layer.eps), average(layer.average), @@ -79,6 +83,7 @@ BatchNormType::BatchNormType(const BatchNormType& layer) : template BatchNormType::BatchNormType(BatchNormType&& layer) : Layer(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::operator=(const BatchNormType& layer) if (&layer != this) { Layer::operator=(layer); + minAxis = layer.minAxis; maxAxis = layer.maxAxis; eps = layer.eps; average = layer.average; @@ -125,6 +131,7 @@ BatchNormType::operator=( if (&layer != this) { Layer::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::Gradient( template void BatchNormType::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::serialize( { ar(cereal::base_class>(this)); + ar(CEREAL_NVP(minAxis)); ar(CEREAL_NVP(maxAxis)); ar(CEREAL_NVP(eps)); ar(CEREAL_NVP(count)); diff --git a/src/mlpack/tests/ann_layer_test.cpp b/src/mlpack/tests/ann_layer_test.cpp index e9c4a3969b..dd957a98d7 100644 --- a/src/mlpack/tests/ann_layer_test.cpp +++ b/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({ 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({ 1, 4, 2 }); module1.ComputeOutputDimensions();