Method setup_data

src/layers/transform/weights.cpp:90–146

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

  // Initialize default weights if none are provided
  if (!this->has_weights()) {
    auto w = std::make_shared<WeightsType>(*this->get_comm());
    auto init = std::make_unique<constant_initializer<DataType>>(DataType(0));
    auto opt = this->m_model->template create_optimizer<TensorDataType>();
    w->set_name(this->get_name() + "_weights");
    w->set_initializer(std::move(init));
    w->set_optimizer(std::move(opt));
    this->add_weights(w);
    this->m_model->add_weights(std::move(w));
  }
  if (this->num_weights() != 1) {
    LBANN_ERROR("attempted to setup ",
                this->get_type(),
                " layer \"",
                this->get_name(),
                "\" ",
                "with an invalid number of weights ",
                "(expected at most 1, ",
                "but found ",
                this->num_weights(),
                ")");
  }

  // Setup weights
  const auto& output_dims_ = this->get_output_dims();
  std::vector<size_t> output_dims(output_dims_.begin(), output_dims_.end());
  auto dist = this->get_activations().DistData();
  dist.rowDist = El::STAR;
  this->get_weights(0).set_dims(output_dims);
  this->get_weights(0).set_matrix_distribution(dist);

  // Initialize freeze state
  if (this->m_frozen) {
    this->get_weights(0).freeze();
  }
  else {
    this->get_weights(0).unfreeze();
  }
  if (this->get_weights(0).is_frozen() != this->m_frozen) {
    LBANN_ERROR((this->m_frozen ? "" : "un"),
                "frozen ",
                "layer \"",
                this->get_name(),
                "\" has ",
                (this->get_weights(0).is_frozen() ? "" : "un"),
                "frozen ",
                "weights \"",
                this->get_weights(0).get_name(),
                "\"");
  }
}