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)
    reduced_keys = ["cxlm_loss", "mlm_loss", "cxlm_ntokens",
      "cxlm_nsentences", "mlm_ntokens", "mlm_nsentences", "cxlm_sample_size",
      "mlm_sample_size", "cxlm_ncorrect", "momentum"]

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

    loss_sum_cxlm = reduced_log["cxlm_loss"]
    loss_sum_mlm = reduced_log["mlm_loss"]
    loss_sum = loss_sum_cxlm + loss_sum_mlm
    cxlm_ntokens = reduced_log["cxlm_ntokens"]
    cxlm_nsentences = reduced_log["cxlm_nsentences"]
    mlm_ntokens = reduced_log["mlm_ntokens"]
    mlm_nsentences = reduced_log["mlm_nsentences"]
    cxlm_sample_size = reduced_log["cxlm_sample_size"]
    mlm_sample_size = reduced_log["mlm_sample_size"]
    sample_size = cxlm_sample_size + mlm_sample_size
    ncorrect = reduced_log["cxlm_ncorrect"]

    eps = 1e-7

    agg_output = {
      'loss': loss_sum / (sample_size + eps) / math.log(2),
      'ntokens': cxlm_ntokens + mlm_ntokens,
      'nsentences': cxlm_nsentences + mlm_nsentences,
      'xlco_loss': loss_sum_cxlm / (cxlm_sample_size + eps) / math.log(2),
      'mlm_loss': loss_sum_mlm / (mlm_sample_size + eps) / math.log(2),
      'xlco_accuracy': 100.0 * ncorrect / (cxlm_nsentences + eps),
      'momentum': reduced_log["momentum"] / len(logging_outputs),
      'xlco_ntokens': cxlm_ntokens,
      'xlco_nsentences': cxlm_nsentences,
      'mlm_ntokens': mlm_ntokens,
      'mlm_nsentences': mlm_nsentences,
      'xlco_sample_size': cxlm_sample_size,
      'mlm_sample_size': mlm_sample_size,
      'sample_size': sample_size,
    }

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

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