Method get_partner_model

src/execution_algorithms/ltfb/checkpoint_binary.cpp:48–76

{}

std::unique_ptr<model>
CheckpointBinary::get_partner_model(model const& m,
                                    El::Int partner_trainer,
                                    size_t /*step*/)
{
  auto const& comm = *m.get_comm();

  // Start by copying this model, then do the exchange.
  auto partner_model_ptr = std::make_unique<model>(m);
  model& partner_model = *partner_model_ptr;

  // Keep track of weights that shouldn't be exchanged
  std::unordered_map<std::string, std::unique_ptr<weights>> restore_weights;
  auto const& weights_names = this->weights_names();
  if (!weights_names.empty()) {
    for (auto w : partner_model.get_weights()) {
      if (weights_names.find(w->get_name()) == weights_names.cend()) {
        using TensorDataType = DataType;
        using WeightsType = data_type_weights<TensorDataType>;
        restore_weights[w->get_name()] =
          std::make_unique<WeightsType>(dynamic_cast<WeightsType&>(*w));
      }
    }
  }
  exchange(comm, partner_model, partner_trainer);
  restore_model_weights(partner_model, restore_weights);

  return partner_model_ptr;
}

} // namespace ltfb
} // namespace lbann