Method setup_data

src/layers/learning/base_convolution.cpp:328–425

template <typename TensorDataType, El::Device Device>
void base_convolution_layer<TensorDataType, Device>::setup_data(
  size_t max_mini_batch_size)
{
  data_type_layer<TensorDataType>::setup_data(max_mini_batch_size);

  // Tensor dimensions
  const auto& input_dims = this->get_input_dims();
  const auto& output_dims = this->get_output_dims();
  const auto& kernel_dims_ = this->get_kernel_dims();
  std::vector<size_t> kernel_dims(kernel_dims_.begin(), kernel_dims_.end());
  const auto kernel_size = get_linear_size(kernel_dims);

  // Initialize default weights if none are provided
  if (this->num_weights() > 2) {
    LBANN_ERROR("attempted to setup layer \"",
                this->get_name(),
                "\" "
                "with an invalid number of weights "
                "(expected at most 2, "
                "found ",
                this->num_weights(),
                ")");
  }
  if (m_bias_scaling_factor != El::TypeTraits<ScalingType>::Zero()) {
    this->set_num_weights(2);
  }
  else {
    this->set_num_weights(1);
  }
  if (!this->has_weights(0)) {
    auto w = std::make_shared<WeightsType>(*this->get_comm());
    auto init = std::make_unique<he_initializer<TensorDataType>>(
      probability_distribution::gaussian);
    auto opt = this->m_model->template create_optimizer<TensorDataType>();

    w->set_name(this->get_name() + "_kernel");
    w->set_initializer(std::move(init));
    w->set_optimizer(std::move(opt));
    this->set_weights(0, w);
    this->m_model->add_weights(std::move(w));
  }
  auto& kernel_weights = this->get_weights(0);

  // Initialize variance scaling initialization
  if (auto* initializer = kernel_weights.get_initializer()) {
    set_fan_in(*initializer, kernel_size / output_dims[0]);
    set_fan_out(*initializer, kernel_size / input_dims[0]);
  }

  // Initialize weight matrices
  auto dist = this->get_prev_activations().DistData();
  dist.colDist = El::STAR;
  dist.rowDist = El::STAR;
  kernel_weights.set_dims(kernel_dims);
  kernel_weights.set_matrix_distribution(dist);

  // Set up bias if needed.
  if (m_bias_scaling_factor != El::TypeTraits<ScalingType>::Zero()) {