Function bp_compute_impl

src/layers/learning/gru.cpp:785–969

template <typename TensorDataType>
void bp_compute_impl(
  gru_layer<TensorDataType, data_layout::DATA_PARALLEL, El::Device::CPU>& l)
{

  // Matrices
  using LocalMat = El::Matrix<TensorDataType, El::Device::CPU>;
  const auto& output_sequence_grad =
    dynamic_cast<const LocalMat&>(l.get_local_prev_error_signals());
  auto& input_sequence_grad =
    dynamic_cast<LocalMat&>(l.get_local_error_signals(0));
  auto& init_hidden_grad =
    dynamic_cast<LocalMat&>(l.get_local_error_signals(1));

  // Dimensions
  const int local_mini_batch_size = output_sequence_grad.Width();
  const int sequence_length = l.get_input_dims(0)[0];
  const int input_size = l.get_input_size(0) / sequence_length;
  const int hidden_size = l.m_hidden_size;
  const int num_layers = l.m_num_layers;

  // oneDNN objects
  if (l.m_onednn_cpu_objects == nullptr) {
    LBANN_ERROR(l.get_type(),
                " layer \"",
                l.get_name(),
                "\" ",
                "attempted to run oneDNN CPU implementation ",
                "before initializing oneDNN objects");
  }
  constexpr auto Device = El::Device::CPU;
  using Backend = onednn_backend<Device>;
  const auto data_type = Backend::template data_type<TensorDataType>();
  auto& onednn_objects = *l.m_onednn_cpu_objects;
  auto sync_info =
    force(El::MakeMultiSync(get_sync_info(input_sequence_grad),
                            get_sync_info(init_hidden_grad),
                            get_sync_info(output_sequence_grad)));
  auto& engine = onednn::get_device_engine<Device>();
  auto stream = onednn::get_stream<Device>(engine, sync_info);

  // Define closure to send weight gradients to optimizers
  std::vector<LocalMat> weights_grad_list(4 * num_layers);
  for (int i = 0; i < num_layers; ++i) {
    weights_grad_list[4 * i].Resize(3 * hidden_size, input_size);
    weights_grad_list[4 * i + 1].Resize(3 * hidden_size, hidden_size);
    weights_grad_list[4 * i + 2].Resize(3 * hidden_size, 1);
    weights_grad_list[4 * i + 3].Resize(3 * hidden_size, 1);
  }
  auto send_weight_grads_to_optimizers = [&]() {
    TensorDataType buf_scale, in_scale;
    for (int i = 0; i < 4 * num_layers; ++i) {
      auto&& opt = l.get_weights(i).get_optimizer();
      if (opt != nullptr) {
        auto& buf = opt->get_gradient_buffer(buf_scale, in_scale, true);
        El::Scale(buf_scale, buf);
        El::Axpy(in_scale, weights_grad_list[i], buf.Matrix());
      }
    }