Method fp_compute

src/layers/regularizers/batch_normalization.cpp:37–175

template <typename TensorDataType, data_layout T_layout, El::Device Dev>
void batch_normalization_layer<TensorDataType, T_layout, Dev>::fp_compute()
{
  const TensorDataType zero = El::TypeTraits<TensorDataType>::Zero();
  const TensorDataType one = El::TypeTraits<TensorDataType>::One();
  const bool is_training =
    this->m_model->get_execution_context().get_execution_mode() ==
    execution_mode::training;

  // Matrices
  const auto& input = this->get_prev_activations();
  const auto& local_input = input.LockedMatrix();
  auto& local_output = this->get_local_activations();

  // Matrix parameters
  const auto& width = input.Width();
  const auto& local_width = local_input.Width();
  const auto& output_dims = this->get_output_dims();
  const auto& num_channels = output_dims[0];
  const auto& channel_size = this->get_output_size() / num_channels;

  const int correction = this->m_bessel_correction ? 1 : 0;

  // Compute statistics
  if (is_training) {
    using ValuesGetter = weights_details::SafeWeightsAccessor<TensorDataType>;
    // Local matrices
    auto& local_mean = this->m_mean_v->Matrix();
    auto& local_var = this->m_var_v->Matrix();
    auto& local_running_mean =
      ValuesGetter::mutable_values(this->get_weights(2)).Matrix();
    auto& local_running_var =
      ValuesGetter::mutable_values(this->get_weights(3)).Matrix();
    // Compute sums and sums of squares
    LBANN_OMP_PARALLEL_FOR
    for (El::Int channel = 0; channel < num_channels; ++channel) {
      TensorDataType sum = zero;
      TensorDataType sqsum = zero;
      const auto& row_start = channel * channel_size;
      const auto& row_end = (channel + 1) * channel_size;
      for (El::Int col = 0; col < local_width; ++col) {
        for (El::Int row = row_start; row < row_end; ++row) {
          const auto& x = local_input(row, col);
          sum += x;
          sqsum += x * x;
        }
      }
      local_mean(channel, 0) = sum;
      local_var(channel, 0) = sqsum;
    }
    El::Int num_per_sum;
    if (this->m_statistics_group_size == 0) {
      // Global statistics aggregation; allreduce on fused buffer.
      this->get_comm()->allreduce(*this->m_mean_and_var,
                                  this->m_mean_and_var->RedundantComm(),
                                  El::mpi::SUM);
      num_per_sum = channel_size * width;
    }
    else if (this->m_statistics_group_size == 1) {