def aggregate_logging_outputs(logging_outputs):
    """Aggregate logging outputs from data parallel training."""
    # loss_sum = sum(log.get('loss', 0) for log in logging_outputs)

    reduced_log = collections.defaultdict(float)
    # TODO sa EM & F1
    reduced_keys = ["sa_loss", 'sa_EM', 'sa_EM_tot', 'sa_nsentences', 'sa_ntokens', 'sa_sample_size', "tlm_loss", "tlm_sample_size", "mlm_ntokens", "mlm_nsentences", "mlm_sample_size", "mlm_loss"]

    for log in logging_outputs:
      for key in reduced_keys:
        reduced_log[key] += log.get(key, 0)

    eps = 1e-7
    sa_sample_size = reduced_log["sa_sample_size"]
    sa_loss = reduced_log["sa_loss"] / (sa_sample_size + eps) / math.log(2)
    tlm_sample_size = reduced_log["tlm_sample_size"]
    tlm_loss = reduced_log["tlm_loss"] / (tlm_sample_size + eps) / math.log(2)
    mlm_sample_size = reduced_log["mlm_sample_size"]
    mlm_loss = reduced_log["mlm_loss"] / (mlm_sample_size + eps) / math.log(2)
    sample_size = sa_sample_size + tlm_sample_size + mlm_sample_size
    loss = (reduced_log["sa_loss"] + reduced_log["tlm_loss"] + reduced_log["mlm_loss"]) / (sample_size + eps) / math.log(2)

    # WALKAROUND
    if reduced_log["sa_EM_tot"] < 1: reduced_log["sa_EM_tot"] = 1

    agg_output = {
      'loss': loss,
      'ntokens': reduced_log["sa_ntokens"] + reduced_log["mlm_ntokens"],
      'nsentences': reduced_log["sa_nsentences"] + reduced_log["mlm_nsentences"],
      'dwa_loss': sa_loss,
      'dwa_sample_size': sa_sample_size,
      'dwa_EM': 0 if reduced_log["sa_EM_tot"] == 0 else 100 * reduced_log["sa_EM"] / reduced_log["sa_EM_tot"],
      'mlm_loss': mlm_loss,
      'mlm_sample_size': mlm_sample_size,
      'tlm_loss': tlm_loss,
      'tlm_sample_size': tlm_sample_size,
      'sample_size': sample_size,
    }

    # DEBUG
    # for k, v in agg_output.items():
    #   print("%s: %.2f" % (k, v), end=" | ")
    # print("")

    return agg_output
Aggregate logging outputs from data parallel training.
infoxlm/src-infoxlm/infoxlm/criterions/xlm_align.py