train.py

# coding: utf-8
###
 # @file   train.py
 # @author Sébastien Rouault <sebastien.rouault@alumni.epfl.ch>
 #
 # @section LICENSE
 #
 # Copyright © 2019-2021 École Polytechnique Fédérale de Lausanne (EPFL).
 # See LICENSE file.
 #
 # @section DESCRIPTION
 #
 # Simulate a training session under attack.
###

import tools
tools.success("Module loading...")

import argparse
import collections
import json
import math
import os
import pathlib
import random
import signal
import sys
import torch
import torchvision
import traceback

import aggregators
import attacks
import experiments

def process_commandline():
  """ Parse the command-line and perform checks.
  Returns:
    Parsed configuration
  """
  # Description
  parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
  parser.add_argument("--seed",
    type=int,
    default=-1,
    help="Fixed seed to use for reproducibility purpose, negative for random seed")
  parser.add_argument("--device",
    type=str,
    default="cpu",
    help="Device on which to run the experiment, \"auto\" by default")
  parser.add_argument("--device-gar",
    type=str,
    default="same",
    help="Device on which to run the GAR, \"same\" for no change of device")
  parser.add_argument("--nb-steps",
    type=int,
    default=300,
    help="Number of training steps to do, non-positive for no limit")
  parser.add_argument("--nb-workers",
    type=int,
    default=11,
    help="Total number of worker machines")
  parser.add_argument("--nb-decl-byz",
    type=int,
    default=4,
    help="Number of Byzantine worker(s) to support")
  parser.add_argument("--nb-real-byz",
    type=int,
    default=0,
    help="Number of actual Byzantine worker(s)")
  parser.add_argument("--gar",
    type=str,
    default="average",
    help="(Byzantine-resilient) aggregation rule to use")
  parser.add_argument("--gar-args",
    nargs="*",
    help="Additional GAR-dependent arguments to pass to the aggregation rule")
  parser.add_argument("--privacy",
    action="store_true",
    default=False,
    help="Gaussian privacy noise ε constant")
  parser.add_argument("--privacy-epsilon",
    type=float,
    default=0.5,
    help="Gaussian privacy noise ε constant; ignore if '--privacy' is not specified")
  parser.add_argument("--privacy-delta",
    type=float,
    default=0.5,
    help="Gaussian privacy noise δ constant; ignore if '--privacy' is not specified")
  parser.add_argument("--gradient-clip",
    type=float,
    default=5,
    help="Maximum L2-norm, above which clipping occurs, for the estimated gradients")
  parser.add_argument("--attack",
    type=str,
    default="nan",
    help="Attack to use")
  parser.add_argument("--attack-args",
    nargs="*",
    help="Additional attack-dependent arguments to pass to the attack")
  parser.add_argument("--model",
    type=str,
    default="simples-conv",
    help="Model to train")
  parser.add_argument("--model-args",
    nargs="*",
    help="Additional model-dependent arguments to pass to the model")
  parser.add_argument("--loss",
    type=str,
    default="nll",
    help="Loss to use")
  parser.add_argument("--loss-args",
    nargs="*",
    help="Additional loss-dependent arguments to pass to the loss")
  parser.add_argument("--criterion",
    type=str,
    default="top-k",
    help="Criterion to use")
  parser.add_argument("--criterion-args",
    nargs="*",
    help="Additional criterion-dependent arguments to pass to the criterion")
  parser.add_argument("--dataset",
    type=str,
    default="mnist",
    help="Dataset to use")
  parser.add_argument("--dataset-args",
    nargs="*",
    help="Additional dataset-dependent arguments to pass to the dataset")
  parser.add_argument("--batch-size",
    type=int,
    default=25,
    help="Batch-size to use for training, 0 for maximum")
  parser.add_argument("--batch-size-test",
    type=int,
    default=100,
    help="Batch-size to use for testing, 0 for maximum")
  parser.add_argument("--test-repeat",
    type=int,
    default=100,
    help="How many evaluation(s) with the test batch-size to average for one evaluation")
  parser.add_argument("--no-transform",
    action="store_true",
    default=False,
    help="Whether to disable any dataset tranformation (e.g. random flips)")
  parser.add_argument("--learning-rate",
    type=float,
    default=0.01,
    help="Learning rate to use for training")
  parser.add_argument("--learning-rate-decay",
    type=int,
    default=5000,
    help="Learning rate hyperbolic half-decay time, non-positive for no decay")
  parser.add_argument("--learning-rate-decay-delta",
    type=int,
    default=5000,
    help="How many steps between two learning rate updates, must be a positive integer")
  parser.add_argument("--momentum",
    type=float,
    default=0.9,
    help="Momentum to use for training")
  parser.add_argument("--dampening",
    type=float,
    default=0.,
    help="Dampening to use for training")
  parser.add_argument("--weight-decay",
    type=float,
    default=0,
    help="Weight decay to use for training")
  parser.add_argument("--l1-regularize",
    type=float,
    default=None,
    help="Add L1 regularization of the given factor to the loss")
  parser.add_argument("--l2-regularize",
    type=float,
    default=None,
    help="Add L2 regularization of the given factor to the loss")
  parser.add_argument("--result-directory",
    type=str,
    default='results-data',
    help="Path of the directory in which to save the experiment results (loss, cross-accuracy, ...) and checkpoints, empty for no saving")
  parser.add_argument("--evaluation-delta",
    type=int,
    default=10,
    help="How many training steps between model evaluations, 0 for no evaluation")
  parser.add_argument("--user-input-delta",
    type=int,
    default=0,
    help="How many training steps between two prompts for user command inputs, 0 for no user input")
  # Parse command line
  return parser.parse_args(sys.argv[1:])

def cmd_make_tree(subtree, level=0):
  if isinstance(subtree, tuple) and len(subtree) > 0 and isinstance(subtree[0], tuple) and len(subtree[0]) == 2:
    label_len = max(len(label) for label, _ in subtree)
    iterator  = subtree
  elif isinstance(subtree, dict):
    if len(subtree) == 0:
      return " - <none>"
    label_len = max(len(label) for label in subtree.keys())
    iterator  = subtree.items()
  else:
    return f" - {subtree}"
  level_spc = "  " * level
  res = ""
  for label, node in iterator:
    res += f"{os.linesep}{level_spc}· {label}{' ' * (label_len - len(label))}{cmd_make_tree(node, level + 1)}"
  return res

def result_make(args, result_fds, name, *fields):
  """ Make and bind a new result file with a name, initialize with a header line.
  Args:
    name      Name of the result file
    fields... Name of each field, in order
  Raises:
    'KeyError' if name is already bound
    'RuntimeError' if no name can be bound
    Any exception that 'io.FileIO' can raise while opening/writing/flushing
  """
  # Check if results are to be output
  
  if args.result_directory is None:
    raise RuntimeError("No result is to be output")
  # Check if name is already bounds
  
  if name in result_fds:
    raise KeyError(f"Name {name!r} is already bound to a result file")
  # Make the new file
  fd = (args.result_directory / name).open("w")
  fd.write("# " + ("\t").join(str(field) for field in fields))
  fd.flush()
  result_fds[name] = fd

def result_get(args, result_fds, name):
  """ Get a valid descriptor to the bound result file, or 'None' if the given name is not bound.
  Args:
    name Given name
  Returns:
    Valid file descriptor, or 'None'
  """
  # Check if results are to be output
  
  if args.result_directory is None:
    return None
  # Return the bound descriptor, if any
  
  return result_fds.get(name, None)

def result_store(fd, *entries):
  """ Store a line in a valid result file.
  Args:
    fd         Descriptor of the valid result file
    entries... Object(s) to convert to string and write in order in a new line
  """
  fd.write(os.linesep + ("\t").join(str(entry) for entry in entries))
  fd.flush()

def compute_avg_dev(values):
  """ Compute the arithmetic mean and standard deviation of a list of values.
  Args:
    values Iterable of values
  Returns:
    Arithmetic mean, standard deviation
  """
  avg = sum(values) / len(values)
  var = 0
  for value in values:
    var += (value - avg) ** 2
  var /= len(values) - 1
  return avg, math.sqrt(var)

class StopTrainingLoop(Exception):
  """ Local exception to signal and stop the training loop.
  """
  pass
  
def main():
  # ---------------------------------------------------------------------------- #
  # Miscellaneous initializations
  tools.success("Miscellaneous initializations...")

  # "Exit requested" global variable accessors
  exit_is_requested, exit_set_requested = tools.onetime("exit")

  # Signal handlers
  signal.signal(signal.SIGINT, exit_set_requested)
  signal.signal(signal.SIGTERM, exit_set_requested)

  # ---------------------------------------------------------------------------- #
  # Command-line processing
  tools.success("Command-line processing...")

  with tools.Context("cmdline", "info"):
    args = process_commandline()
    # Parse additional arguments
    for name in ("gar", "attack", "model", "dataset", "loss", "criterion"):
      name = f"{name}_args"
      keyval = getattr(args, name)
      setattr(args, name, dict() if keyval is None else tools.parse_keyval(keyval))
    # Count the number of real honest workers
    args.nb_honests = args.nb_workers - args.nb_real_byz
    if args.nb_honests < 0:
      tools.fatal(f"Invalid arguments: there are more real Byzantine workers ({args.nb_real_byz}) than total workers ({args.nb_workers})")
    # Check general training parameters
    if args.momentum < 0.:
      tools.fatal(f"Invalid arguments: negative momentum factor {args.momentum}")
    if args.dampening < 0.:
      tools.fatal(f"Invalid arguments: negative dampening factor {args.dampening}")
    if args.weight_decay < 0.:
      tools.fatal(f"Invalid arguments: negative weight decay factor {args.weight_decay}")
    # Check the learning rate and associated options
    if args.learning_rate <= 0:
      tools.fatal(f"Invalid arguments: non-positive learning rate {args.learning_rate}")
    if args.learning_rate_decay < 0:
      tools.fatal(f"Invalid arguments: negative learning rate decay {args.learning_rate_decay}")
    if args.learning_rate_decay_delta <= 0:
      tools.fatal(f"Invalid arguments: non-positive learning rate decay delta {args.learning_rate_decay_delta}")
    # Check the privacy-related metrics
    if args.gradient_clip <= 0.:
      tools.fatal(f"Invalid arguments: non-positive gradient clip constant {args.gradient_clip}")
    if args.privacy:
      if args.privacy_epsilon <= 0. or args.privacy_epsilon >= 1.:
        tools.fatal(f"Invalid arguments: off-bounds (]0, 1[) ε constant {args.privacy_epsilon}")
      if args.privacy_delta <= 0. or args.privacy_delta >= 1.:
        tools.fatal(f"Invalid arguments: off-bounds (]0, 1[) δ constant {args.privacy_delta}")
      args.privacy_sensitivity = 2 * args.gradient_clip / args.batch_size
    # Print configuration

    cmdline_config = "Configuration" + cmd_make_tree((
      ("Reproducibility", "not enforced" if args.seed < 0 else f"enforced (seed {args.seed})"),
      ("#workers", args.nb_workers),
      ("#declared Byz.", args.nb_decl_byz),
      ("#actually Byz.", args.nb_real_byz),
      ("Model", (
        ("Name", args.model),
        ("Arguments", args.model_args))),
      ("Dataset", (
        ("Name", args.dataset),
        ("Arguments", args.dataset_args),
        ("Batch size", (
          ("Training", args.batch_size or 'max'),
          ("Testing", f"{args.batch_size_test or 'max'} × {args.test_repeat}"))),
        ("Transforms", "none" if args.no_transform else "default"))),
      ("Loss", (
        ("Name", args.loss),
        ("Arguments", args.loss_args),
        ("Regularization", (
          ("l1", "none" if args.l1_regularize is None else args.l1_regularize),
          ("l2", "none" if args.l2_regularize is None else args.l2_regularize))))),
      ("Criterion", (
        ("Name", args.criterion),
        ("Arguments", args.criterion_args))),
      ("Optimizer", (
        ("Name", "sgd"),
        ("Learning rate", (
          ("Initial", args.learning_rate),
          ("Half-decay", args.learning_rate_decay if args.learning_rate_decay > 0 else "none"),
          ("Update delta", args.learning_rate_decay_delta if args.learning_rate_decay > 0 else "n/a"))),
        ("Momentum", args.momentum),
        ("Dampening", args.dampening),
        ("Weight decay", args.weight_decay))),
      ("Attack", (
        ("Name", args.attack),
        ("Arguments", args.attack_args))),
      ("Aggregation", (
        ("Name", args.gar),
        ("Arguments", args.gar_args))),
      ("Differential privacy", (
        ("Enabled?", "yes" if args.privacy else "no"),
        ("ε constant", args.privacy_epsilon if args.privacy else "n/a"),
        ("δ constant", args.privacy_delta if args.privacy else "n/a"),
        ("l2-sensitivity", args.privacy_sensitivity if args.privacy else "n/a")))))
    print(cmdline_config)

  # ---------------------------------------------------------------------------- #
  # Setup
  tools.success("Experiment setup...")

  with tools.Context("setup", "info"):
    # Enforce reproducibility if asked (see https://pytorch.org/docs/stable/notes/randomness.html)
    reproducible = args.seed >= 0
    if reproducible:
      torch.manual_seed(args.seed)
      import numpy
      numpy.random.seed(args.seed)
    torch.backends.cudnn.deterministic = reproducible
    torch.backends.cudnn.benchmark     = not reproducible
    # Configurations
    config = experiments.Configuration(dtype=torch.float32, device=(None if args.device.lower() == "auto" else args.device), noblock=True)
    if args.device_gar.lower() == "same":
      config_gar = config
    else:
      config_gar = experiments.Configuration(dtype=config["dtype"], device=(None if args.device_gar.lower() == "auto" else args.device_gar), noblock=config["non_blocking"])
    # Defense
    defense = aggregators.gars.get(args.gar)
    if defense is None:
      tools.fatal_unavailable(aggregators.gars, args.gar, what="aggregation rule")
    # Attack
    attack = attacks.attacks.get(args.attack)
    if attack is None:
      tools.fatal_unavailable(attacks.attacks, args.attack, what="attack")
    # Model
    model = experiments.Model(args.model, config, **args.model_args)

    # Datasets
    if args.no_transform:
      train_transforms = test_transforms = torchvision.transforms.ToTensor()
    else:
      train_transforms = test_transforms = None # Let default values
    trainset, testset = experiments.make_datasets(args.dataset, args.batch_size, args.batch_size_test, train_transforms=train_transforms, test_transforms=test_transforms, **args.dataset_args)
    model.default("trainset", trainset)
    model.default("testset", testset)
    # Loss and criterion
    loss = experiments.Loss(args.loss, **args.loss_args)
    if args.l1_regularize is not None:
      loss += args.l1_regularize * experiments.Loss("l1")
    if args.l2_regularize is not None:
      loss += args.l2_regularize * experiments.Loss("l2")
    criterion = experiments.Criterion(args.criterion, **args.criterion_args)
    model.default("loss", loss)
    model.default("criterion", criterion)
    # Optimizer
    optimizer = experiments.Optimizer("sgd", model, lr=args.learning_rate, momentum=args.momentum, dampening=args.dampening, weight_decay=args.weight_decay)
    model.default("optimizer", optimizer)
    # Privacy noise distribution
    if args.privacy:
      param = model.get()
      privacy_factor = args.privacy_sensitivity * math.sqrt(2 * math.log(1.25 / args.privacy_delta)) / args.privacy_epsilon
      grad_noise = torch.distributions.normal.Normal(torch.zeros_like(param), torch.ones_like(param).mul_(privacy_factor))
    # Make the result directory (if requested)
    if args.result_directory is not None:
      try:
        resdir = pathlib.Path(args.result_directory).resolve()
        resdir.mkdir(mode=0o755, parents=True, exist_ok=True)
        args.result_directory = resdir
      except Exception as err:
        tools.warning(f"Unable to create the result directory {str(resdir)!r} ({err}); no result will be stored")
      else:
        result_fds = dict()
        try:
          # Make evaluation file
          if args.evaluation_delta > 0:
            result_make(args, result_fds, "eval", "Step number", "Cross-accuracy")
          # Make study file
          result_make(args, result_fds, "study", "Step number", "Training point count",
            "Average loss", "l2 from origin",
            "Honest gradient deviation", "Attack gradient deviation",
            "Honest gradient norm", "Attack gradient norm", "Defense gradient norm",
            "Honest max coordinate", "Attack max coordinate", "Defense max coordinate",
            "Honest-attack cosine", "Honest-defense cosine", "Attack-defense cosine")
          # Store the configuration info and JSON representation
          (args.result_directory / "config").write_text(cmdline_config + os.linesep)
          with (args.result_directory / "config.json").open("w") as fd:
            def convert_to_supported_json_type(x):
              if type(x) in {str, int, float, bool, type(None), dict, list}:
                return x
              elif type(x) is set:
                return list(x)
              else:
                return str(x)
            datargs = dict((name, convert_to_supported_json_type(getattr(args, name))) for name in dir(args) if len(name) > 0 and name[0] != "_")
            del convert_to_supported_json_type
            json.dump(datargs, fd, ensure_ascii=False, indent="\t")
        except Exception as err:
          tools.warning(f"Unable to create some result files in directory {str(resdir)!r} ({err}); some result(s) may be missing")

  # ---------------------------------------------------------------------------- #
  # Training
  tools.success("Training...")


  # Training until limit or stopped
  with tools.Context("training", "info"):
    was_training  = False
    current_lr    = args.learning_rate
    steps         = 0
    datapoints    = 0
    fd_eval       = result_get(args, result_fds, "eval")
    fd_study      = result_get(args, result_fds, "study")
    atc_gradient  = tools.AccumulatedTimedContext(sync=True)
    atc_noise     = tools.AccumulatedTimedContext(sync=True)
    atc_aggregate = tools.AccumulatedTimedContext(sync=True)
    atc_evaluate  = tools.AccumulatedTimedContext(sync=True)
    params_origin = model.get().clone().detach_()
    try:
      while not exit_is_requested():
        # ------------------------------------------------------------------------ #
        # Evaluate if any milestone is reached
        milestone_evaluation = args.evaluation_delta > 0 and steps % args.evaluation_delta == 0
        milestone_user_input = args.user_input_delta > 0 and steps % args.user_input_delta == 0
        milestone_any        = milestone_evaluation or milestone_user_input
        # Training notification (end)
        if milestone_any and was_training:
          print(" done.")
          was_training = False
        # Evaluation milestone reached
        if milestone_evaluation:
          print("Accuracy (step %d)..." % steps, end="", flush=True)
          with atc_evaluate:
            res = model.eval()
            for _ in range(args.test_repeat - 1):
              res += model.eval()
            acc = res[0].item() / res[1].item()
          print(" %.2f%%." % (acc * 100.))
          # Store the evaluation result
          if fd_eval is not None:
            result_store(fd_eval, steps, acc)
        # User input milestone
        if milestone_user_input:
          tools.interactive()
        # Check if reach step limit
        if args.nb_steps > 0 and steps >= args.nb_steps:
          # Training notification (end)
          if was_training:
            print(" done.")
            was_training = False
          # Leave training loop
          raise StopTrainingLoop()
        # Training notification (begin)
        if milestone_any and not was_training:
          print("Training...", end="", flush=True)
          was_training = True
        # ------------------------------------------------------------------------ #
        # Compute (honest) losses (if it makes sense), gradients and voting data
        grad_honests = list()
        loss_honests = list()
        # For each honest worker
        with atc_gradient:
          for i in range(args.nb_honests):
            grad, loss = model.backprop(outloss=True)
            grad = grad.clone().detach_()
            # Loss append
            loss_honests.append(loss.item())
            # Gradient clip
            if args.gradient_clip is not None:
              grad_norm = grad.norm().item()
              if grad_norm > args.gradient_clip:
                grad.mul_(args.gradient_clip / grad_norm)
            # Gradient append
            grad_honests.append(grad)
        # Pre-compute some quantities for study (here if necessary, since each gradient in 'grad_honests' may have privacy noise added)
        if fd_study is not None:
          # Compute average loss ('len(loss_honests) > 0' is guaranteed)
          loss_avg = sum(loss_honests) / len(loss_honests)
          # Compute the sampled and honest gradients norm average, norm deviation and max absolute coordinate
          honest_grad_avg, honest_norm_avg, honest_norm_dev, honest_norm_max = tools.compute_avg_dev_max(grad_honests)
        # Move the honest gradients to the GAR device
        if config_gar is not config:
          grad_honests_gar = list(grad.to(device=config_gar["device"], non_blocking=config_gar["non_blocking"]) for grad in grad_honests)
        else:
          grad_honests_gar = grad_honests
        # ------------------------------------------------------------------------ #
        # Add privacy noise to the 'grad_honests_gar' (might be 'grad_honests'), which are the ones sent and processed by the server
        if args.privacy:
          with atc_noise:
            for grad in grad_honests_gar:
              # QC: change
              if args.privacy_epsilon == 0.5:
                print("======== sparsity triggered =======")
                grad_positive = grad[grad > 0]
                grad_negative = grad[grad < 0]
                if len(grad_positive) > 0:
                  grad_p = torch.quantile(grad_positive, 0.25)
                  grad[(grad > 0) * (grad < grad_p)] = 0
                if len(grad_negative) > 0:
                  grad_n = torch.quantile(grad_negative, 0.75)
                  grad[(grad < 0) * (grad > grad_n)] = 0
              else:
                grad.add_(grad_noise.sample())

        # ------------------------------------------------------------------------ #
        # Compute the Byzantine gradients
        grad_attacks = attack.checked(grad_honests=grad_honests_gar, f_decl=args.nb_decl_byz, f_real=args.nb_real_byz, model=model, defense=defense, **args.attack_args)
        # ------------------------------------------------------------------------ #
        # Aggregate and update the model
        with atc_aggregate:
          grad_defense = defense.checked(gradients=(grad_honests_gar + grad_attacks), f=args.nb_decl_byz, model=model, **args.gar_args)
        # Move the defense gradient back to the main device
        if config_gar is not config:
          for grad in grad_attacks:
            grad.data = grad.to(device=config["device"], non_blocking=config["non_blocking"])
          grad_defense = grad_defense.to(device=config["device"], non_blocking=config["non_blocking"])
        # Compute l2-distance from origin (if needed for study)
        if fd_study is not None:
          l2_origin = model.get().sub(params_origin).norm().item()
        # Model update (possibly updating the learning rate)
        if args.learning_rate_decay > 0 and steps % args.learning_rate_decay_delta == 0:
          current_lr = args.learning_rate / (steps / args.learning_rate_decay + 1)
          optimizer.set_lr(current_lr)
        model.update(grad_defense)
        # ------------------------------------------------------------------------ #
        # Store study (if requested)
        if fd_study is not None:
          # Compute the sampled and honest gradients norm average, norm deviation and max absolute coordinate
          attack_grad_avg, attack_norm_avg, attack_norm_dev, attack_norm_max = tools.compute_avg_dev_max(grad_attacks)
          # Compute the defense norm average and max absolute coordinate
          defense_grad = grad_defense # (Mere renaming for consistency)
          defense_norm_avg = defense_grad.norm().item()
          defense_norm_max = defense_grad.abs().max().item()
          # Compute cosine of solid angles
          cosin_honatt = math.nan if attack_grad_avg is None else torch.dot(honest_grad_avg, attack_grad_avg).div_(honest_norm_avg).div_(attack_norm_avg).item()
          cosin_hondef = torch.dot(honest_grad_avg, defense_grad).div_(honest_norm_avg).div_(defense_norm_avg).item()
          cosin_attdef = math.nan if attack_grad_avg is None else torch.dot(attack_grad_avg, defense_grad).div_(attack_norm_avg).div_(defense_norm_avg).item()
          # Store the result (float-to-string format chosen so not to lose precision)
          float_format = {torch.float16: "%.4e", torch.float32: "%.8e", torch.float64: "%.16e"}.get(config["dtype"], "%s")
          result_store(fd_study, steps, datapoints,
            float_format % loss_avg, float_format % l2_origin,
            float_format % honest_norm_dev, float_format % attack_norm_dev,
            float_format % honest_norm_avg, float_format % attack_norm_avg, float_format % defense_norm_avg,
            float_format % honest_norm_max, float_format % attack_norm_max, float_format % defense_norm_max,
            float_format % cosin_honatt, float_format % cosin_hondef, float_format % cosin_attdef)
        # ------------------------------------------------------------------------ #
        # Increase the step counter
        steps      += 1
        datapoints += args.batch_size * args.nb_honests
    except StopTrainingLoop:
      pass
    # Training notification (end)
    if was_training:
      print(" interrupted.")

  # Print and store timing counters
  with tools.Context("perf", "info"):
    perfs = dict()
    perf_params = (
      (atc_gradient,  "grad",  "Gradient computation (per worker)", args.nb_honests),
      (atc_noise,     "noise", "Noise addition (per worker)",       args.nb_honests),
      (atc_aggregate, "aggr",  "Gradient aggregation", 1),
      (atc_evaluate,  "eval",  "Model evaluation",     1))
    # Compute max name length
    nlen = max(len(name) for _, _, name, _ in perf_params)
    # Print
    for atc, key, name, div in perf_params:
      acc = tools.AccumulatedTimedContext(atc.current_runtime() / div)
      print(f"{name:{nlen}s} - {acc}")
      perfs[key] = (acc.current_runtime(), name)
    # Store
    if args.result_directory:
      with (args.result_directory / "perfs.json").open("w") as fd:
        json.dump(perfs, fd, ensure_ascii=False, indent="\t")

if __name__ == '__main__':
    main()