diff --git a/src/mlpack/methods/ann/layer/batch_norm.hpp b/src/mlpack/methods/ann/layer/batch_norm.hpp index 8c07fd0fc5..7b20a9c27e 100644 --- a/src/mlpack/methods/ann/layer/batch_norm.hpp +++ b/src/mlpack/methods/ann/layer/batch_norm.hpp @@ -64,7 +64,7 @@ class BatchNorm * * @param size The number of input units / channels. * @param eps The epsilon added to variance to ensure numerical stability. - * @param average Boolean to determine whether counting average is used for + * @param average Boolean to determine whether cumulative average is used for * updating the parameters or momentum is used. * @param momentum Parameter used to to update the running mean and variance. */ diff --git a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp index 11f5ff8a9a..0d195603ac 100644 --- a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp +++ b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp @@ -91,6 +91,14 @@ void BatchNorm::Forward( // We will calculate minibatch norm on each channel / feature map. if (!deterministic) { + // Check only during training, batch-size can be one during inference. + if (batchSize == 1 && inputSize == 1) + { + Log::Warn << "Variance for single element isn't defined and" << + " will be set to 0.0 for training. Use a batch-size" << + " greater than 1 to fix the warning." << std::endl; + } + // Input corresponds to output from convolution layer. // Use a cube for simplicity. arma::cube inputTemp(const_cast&>(input).memptr(), @@ -129,11 +137,15 @@ void BatchNorm::Forward( count += 1; averageFactor = average ? 1.0 / count : momentum; + double nElements = 0.0; + if (input.n_elem - size != 0) + nElements = 1.0 / (input.n_elem - size + eps); + // Update running mean and running variance. runningMean = (1 - averageFactor) * runningMean + averageFactor * mean.t(); runningVariance = (1 - averageFactor) * runningVariance + - input.n_elem * (1.0 / (input.n_elem - size + eps)) * + input.n_elem * nElements * averageFactor * variance.t(); } else