code: # Copyright 2020, Google LLC.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Runs federated training on various tasks using a generalized form of FedAvg.
  
  Specifically, we create (according to flags) an iterative processes that allows
  for client and server learning rate schedules, as well as various client and
  server optimization methods. For more details on the learning rate scheduling
  and optimization methods, see `shared/optimizer_utils.py`. For details on the
  iterative process, see `shared/fed_avg_schedule.py`.
  """
  
  import collections
  import os.path
  from typing import Callable
  
  from absl import app
  from absl import flags
  import tensorflow as tf
  import tensorflow_federated as tff
  
  from optimization.cifar100 import federated_cifar100
  from optimization.emnist import federated_emnist
  from optimization.emnist_ae import federated_emnist_ae
  from optimization.shakespeare import federated_shakespeare
  from optimization.shared import fed_avg_schedule
  from optimization.shared import optimizer_utils
  from optimization.shared import training_specs
  from optimization.stackoverflow import federated_stackoverflow
  from optimization.stackoverflow_lr import federated_stackoverflow_lr
  from utils import training_loop
  from utils import utils_impl
  
  _SUPPORTED_TASKS = [
      'cifar100', 'emnist_cr', 'emnist_ae', 'shakespeare', 'stackoverflow_nwp',
      'stackoverflow_lr'
  ]
  
  with utils_impl.record_hparam_flags() as optimizer_flags:
    # Defining optimizer flags
    optimizer_utils.define_optimizer_flags('client')
    optimizer_utils.define_optimizer_flags('server')
    optimizer_utils.define_lr_schedule_flags('client')
    optimizer_utils.define_lr_schedule_flags('server')
  
  with utils_impl.record_hparam_flags() as shared_flags:
    # Federated training hyperparameters
    flags.DEFINE_integer('client_epochs_per_round', 1,
                         'Number of epochs in the client to take per round.')
    flags.DEFINE_integer('client_batch_size', 20, 'Batch size on the clients.')
    flags.DEFINE_integer('clients_per_round', 10,
                         'How many clients to sample per round.')
    flags.DEFINE_integer('client_datasets_random_seed', 1,
                         'Random seed for client sampling.')
  
    # Training loop configuration
    flags.DEFINE_string(
        'experiment_name', None, 'The name of this experiment. Will be append to '
        '--root_output_dir to separate experiment results.')
    flags.mark_flag_as_required('experiment_name')
    flags.DEFINE_string('root_output_dir', '/tmp/fed_opt/',
                        'Root directory for writing experiment output.')
    flags.DEFINE_integer('total_rounds', 200, 'Number of total training rounds.')
    flags.DEFINE_integer(
        'rounds_per_eval', 1,
        'How often to evaluate the global model on the validation dataset.')
    flags.DEFINE_integer('rounds_per_checkpoint', 50,
                         'How often to checkpoint the global model.')
  
  with utils_impl.record_hparam_flags() as task_flags:
    # Task specification
    flags.DEFINE_enum('task', None, _SUPPORTED_TASKS,
                      'Which task to perform federated training on.')
  
  with utils_impl.record_hparam_flags() as cifar100_flags:
    # CIFAR-100 flags
    flags.DEFINE_integer('cifar100_crop_size', 24, 'The height and width of '
                         'images after preprocessing.')
    flags.DEFINE_bool(
        'cifar100_distort_train_images', True, 'If set to True, '
        'train images will be randomly cropped. Otherwise, all '
        'images will simply be resized.')
  
  with utils_impl.record_hparam_flags() as emnist_cr_flags:
    # EMNIST CR flags
    flags.DEFINE_enum(
        'emnist_cr_model', 'cnn', ['cnn', '2nn'], 'Which model to '
        'use. This can be a convolutional model (cnn) or a two '
        'hidden-layer densely connected network (2nn).')
  
  with utils_impl.record_hparam_flags() as shakespeare_flags:
    # Shakespeare flags
    flags.DEFINE_integer(
        'shakespeare_sequence_length', 80,
        'Length of character sequences to use for the RNN model.')
  
  with utils_impl.record_hparam_flags() as so_nwp_flags:
    # Stack Overflow NWP flags
    flags.DEFINE_integer('so_nwp_vocab_size', 10000, 'Size of vocab to use.')
    flags.DEFINE_integer('so_nwp_num_oov_buckets', 1,
                         'Number of out of vocabulary buckets.')
    flags.DEFINE_integer('so_nwp_sequence_length', 20,
                         'Max sequence length to use.')
    flags.DEFINE_integer('so_nwp_max_elements_per_user', 1000, 'Max number of '
                         'training sentences to use per user.')
    flags.DEFINE_integer(
        'so_nwp_num_validation_examples', 10000, 'Number of examples '
        'to use from test set for per-round validation.')
  
  with utils_impl.record_hparam_flags() as so_lr_flags:
    # Stack Overflow LR flags
    flags.DEFINE_integer('so_lr_vocab_tokens_size', 10000,
                         'Vocab tokens size used.')
    flags.DEFINE_integer('so_lr_vocab_tags_size', 500, 'Vocab tags size used.')
    flags.DEFINE_integer(
        'so_lr_num_validation_examples', 10000, 'Number of examples '
        'to use from test set for per-round validation.')
    flags.DEFINE_integer('so_lr_max_elements_per_user', 1000,
                         'Max number of training '
                         'sentences to use per user.')
  
  FLAGS = flags.FLAGS
  
  TASK_FLAGS = collections.OrderedDict(
      cifar100=cifar100_flags,
      emnist_cr=emnist_cr_flags,
      shakespeare=shakespeare_flags,
      stackoverflow_nwp=so_nwp_flags,
      stackoverflow_lr=so_lr_flags)
  
  
  def _write_hparam_flags():
    """Creates an ordered dictionary of hyperparameter flags and writes to CSV."""
    hparam_dict = utils_impl.lookup_flag_values(shared_flags)
  
    # Update with optimizer flags corresponding to the chosen optimizers.
    opt_flag_dict = utils_impl.lookup_flag_values(optimizer_flags)
    opt_flag_dict = optimizer_utils.remove_unused_flags('client', opt_flag_dict)
    opt_flag_dict = optimizer_utils.remove_unused_flags('server', opt_flag_dict)
    hparam_dict.update(opt_flag_dict)
  
    # Update with task-specific flags.
    task_name = FLAGS.task
    if task_name in TASK_FLAGS:
      task_hparam_dict = utils_impl.lookup_flag_values(TASK_FLAGS[task_name])
      hparam_dict.update(task_hparam_dict)
  
    results_dir = os.path.join(FLAGS.root_output_dir, 'results',
                               FLAGS.experiment_name)
    utils_impl.create_directory_if_not_exists(results_dir)
    hparam_file = os.path.join(results_dir, 'hparams.csv')
    utils_impl.atomic_write_series_to_csv(hparam_dict, hparam_file)
  
  
  def main(argv):
    if len(argv) > 1:
      raise app.UsageError('Expected no command-line arguments, '
                           'got: {}'.format(argv))
  
    client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('client')
    server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('server')
  
    client_lr_schedule = optimizer_utils.create_lr_schedule_from_flags('client')
    server_lr_schedule = optimizer_utils.create_lr_schedule_from_flags('server')
  
    def iterative_process_builder(
        model_fn: Callable[[],
                           tff.learning.Model]) -> tff.templates.IterativeProcess:
      """Creates an iterative process using a given TFF `model_fn`.
  
      Args:
        model_fn: A no-arg function returning a `tff.learning.Model`.
  
      Returns:
        A `tff.templates.IterativeProcess`.
      """
      if FLAGS.task == 'shakespeare' or FLAGS.task == 'stackoverflow_nwp':
  
        def client_weight_fn(local_outputs):
          return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
      else:
        client_weight_fn = None
  
      return fed_avg_schedule.build_fed_avg_process(
          model_fn=model_fn,
          client_optimizer_fn=client_optimizer_fn,
          client_lr=client_lr_schedule,
          server_optimizer_fn=server_optimizer_fn,
          server_lr=server_lr_schedule,
          client_weight_fn=client_weight_fn)
  
    task_spec = training_specs.TaskSpec(
        iterative_process_builder=iterative_process_builder,
        client_epochs_per_round=FLAGS.client_epochs_per_round,
        client_batch_size=FLAGS.client_batch_size,
        clients_per_round=FLAGS.clients_per_round,
        client_datasets_random_seed=FLAGS.client_datasets_random_seed)
  
    if FLAGS.task == 'cifar100':
      runner_spec = federated_cifar100.configure_training(
          task_spec,
          crop_size=FLAGS.cifar100_crop_size,
          distort_train_images=FLAGS.cifar100_distort_train_images)
    elif FLAGS.task == 'emnist_cr':
      runner_spec = federated_emnist.configure_training(
          task_spec, model=FLAGS.emnist_cr_model)
    elif FLAGS.task == 'emnist_ae':
      runner_spec = federated_emnist_ae.configure_training(task_spec)
    elif FLAGS.task == 'shakespeare':
      runner_spec = federated_shakespeare.configure_training(
          task_spec, sequence_length=FLAGS.shakespeare_sequence_length)
    elif FLAGS.task == 'stackoverflow_nwp':
      runner_spec = federated_stackoverflow.configure_training(
          task_spec,
          vocab_size=FLAGS.so_nwp_vocab_size,
          num_oov_buckets=FLAGS.so_nwp_num_oov_buckets,
          sequence_length=FLAGS.so_nwp_sequence_length,
          max_elements_per_user=FLAGS.so_nwp_max_elements_per_user,
          num_validation_examples=FLAGS.so_nwp_num_validation_examples)
    elif FLAGS.task == 'stackoverflow_lr':
      runner_spec = federated_stackoverflow_lr.configure_training(
          task_spec,
          vocab_tokens_size=FLAGS.so_lr_vocab_tokens_size,
          vocab_tags_size=FLAGS.so_lr_vocab_tags_size,
          max_elements_per_user=FLAGS.so_lr_max_elements_per_user,
          num_validation_examples=FLAGS.so_lr_num_validation_examples)
    else:
      raise ValueError(
          '--task flag {} is not supported, must be one of {}.'.format(
              FLAGS.task, _SUPPORTED_TASKS))
  
    _write_hparam_flags()
  
    training_loop.run(
        iterative_process=runner_spec.iterative_process,
        client_datasets_fn=runner_spec.client_datasets_fn,
        validation_fn=runner_spec.validation_fn,
        test_fn=runner_spec.test_fn,
        total_rounds=FLAGS.total_rounds,
        experiment_name=FLAGS.experiment_name,
        root_output_dir=FLAGS.root_output_dir,
        rounds_per_eval=FLAGS.rounds_per_eval,
        rounds_per_checkpoint=FLAGS.rounds_per_checkpoint)
  
  
  if __name__ == '__main__':
    app.run(main)
   # Copyright 2019, Google LLC.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Utility for filtering (via class. accuracy) the Federated EMNIST dataset."""
  
  import csv
  import functools
  import os.path
  
  import tensorflow as tf
  import tensorflow_federated as tff
  
  from gans.experiments.emnist import emnist_data_utils
  
  BASE_URL = 'https://storage.googleapis.com/tff-experiments-public/'
  CSVS_BASE_PATH = 'gans/csvs/'
  
  
  @functools.lru_cache(maxsize=1)
  def get_unfiltered_client_data_for_training(batch_size):
    r"""Returns `tff.simulation.datasets.ClientData` of unfiltered Federated EMNIST data.
  
    The data returned will neither be filtered by user nor by example, so training
    can take place with all users and all examples for each user.
  
    Args:
      batch_size: Batch size of output dataset. If None, don't batch.
  
    Returns:
      A tff.simulation.datasets.ClientData` of real images of numbers/letters. The
      data has
      not been filtered.
    """
    return get_filtered_client_data_for_training(None, None, batch_size)
  
  
  @functools.lru_cache(maxsize=1)
  def get_filtered_by_user_client_data_for_training(invert_imagery_probability,
                                                    accuracy_threshold,
                                                    batch_size,
                                                    cache_dir=None):
    r"""Returns `tff.simulation.datasets.ClientData` of filtered Federated EMNIST data.
  
    Input data gets filtered on a per-user basis; users get selected via the
    `accuracy_threshold` criterion, and then training can take place with all
    examples from only the selected users.
  
    Args:
      invert_imagery_probability: The probability that a user\'s image data has
        pixel intensity inverted. E.g., `0p1` corresponds to 0.1, or a 10%
        probability that a user\'s data is flipped. Note that to save time in
        experiment execution, this is precomputed via the ./filter_users.py
        script, and the selection here controls which file to read from.
      accuracy_threshold: Indicates the classification threshold by which a user
        is included in the training population.  E.g., `lt0p882` means any user
        who\'s data cumulatively classifies with <0.882 accuracy would be used for
        training; `gt0p939` means any user who\'s data cumulatively classifies
        with >0.939 accuracy would be used for training. To save time in
        experiment execution, this assignment is precomputed via the
        ./filter_users.py script, and the flag selection here is to indicate which
        file to read from.
      batch_size: Batch size of output dataset. If None, don't batch.
      cache_dir: (Optional) base directory to cache the downloaded files. If None,
        caches in Keras' default cache directory.
  
    Returns:
      A tff.simulation.datasets.ClientData` of real images of numbers/letters. The
      data has
      been filtered by user classification accuracy as per the input arguments.
    """
    path_to_data = os.path.join(CSVS_BASE_PATH,
                                'inv_prob_{}'.format(invert_imagery_probability),
                                'filter_by_user',
                                'acc_{}'.format(accuracy_threshold))
  
    try:
      filename = 'client_ids.csv'
      path_to_read_inversions_csv = tf.keras.utils.get_file(
          fname=filename,
          cache_subdir=path_to_data,
          cache_dir=cache_dir,
          origin=os.path.join(BASE_URL, path_to_data, filename))
    except Exception:
      msg = ('A URL fetch failure was encountered when trying to retrieve '
             'filter-by-user generated csv file with invert_imagery_probability '
             '`{}` and accuracy_threshold `{}`. Please run the ./filter_users.py '
             'script to generate the missing data, and use the `cache_dir` '
             'argument to this method to specify the location of the generated '
             'data csv file.'.format(invert_imagery_probability,
                                     accuracy_threshold))
      raise ValueError(msg)
  
    return get_filtered_client_data_for_training(path_to_read_inversions_csv,
                                                 None, batch_size)
  
  
  @functools.lru_cache(maxsize=1)
  def get_filtered_by_example_client_data_for_training(invert_imagery_probability,
                                                       min_num_examples,
                                                       example_class_selection,
                                                       batch_size,
                                                       cache_dir=None):
    r"""Returns `tff.simulation.datasets.ClientData` of filtered Federated EMNIST data.
  
    Input data gets filtered on a per-example basis. Any user meeting the
    `min_num_examples` criterion is included. The examples are limited to those
    that classified according to the `example_class_selection` criterion.
  
    Args:
      invert_imagery_probability: The probability that a user\'s image data has
        pixel intensity inverted. E.g., `0p1` corresponds to 0.1, or a 10%
        probability that a user\'s data is flipped. Note that to save time in
        experiment execution, this is precomputed via the ./filter_examples.py
        scripts, and the selection here controls which file to read from.
      min_num_examples: Indicates the minimum number of examples that are either
        correct or incorrect (as set by the `example_class_selection` argument) in
        a client\'s local dataset for that client to be considered as part of
        training sub-population. To save time in experiment execution, this
        assignment is precomputed via the ./filter_examples.py script, and the
        flag selection here is to indicate which file to read from.
      example_class_selection: Indicates whether to train on a client\'s correct
        or incorrect examples. To save time in experiment execution, this
        assignment is precomputed via the ./filter_examples.py script, and the
        flag selection here is to indicate which file to read from.
      batch_size: Batch size of output dataset. If None, don't batch.
      cache_dir: (Optional) base directory to cache the downloaded files. If None,
        caches in Keras' default cache directory.
  
    Returns:
      A `tff.simulation.datasets.ClientData` of real images of numbers/letters.
      The data
      has been filtered as per the input arguments (either not filtered, filtered
      by user classification accuracy, or filtered by example classification
      correctness).
    """
    path_to_data = os.path.join(CSVS_BASE_PATH,
                                'inv_prob_{}'.format(invert_imagery_probability),
                                'filter_by_example',
                                'min_num_examples_{}'.format(min_num_examples),
                                '{}'.format(example_class_selection))
  
    try:
      filename = 'client_ids.csv'
      path_to_read_inversions_csv = tf.keras.utils.get_file(
          fname=filename,
          cache_subdir=path_to_data,
          cache_dir=cache_dir,
          origin=os.path.join(BASE_URL, path_to_data, filename))
  
      filename = 'example_indices_map.csv'
      path_to_read_example_indices_csv = tf.keras.utils.get_file(
          fname=filename,
          cache_subdir=path_to_data,
          cache_dir=cache_dir,
          origin=os.path.join(BASE_URL, path_to_data, filename))
    except Exception:
      msg = ('A URL fetch failure was encountered when trying to retrieve '
             'filter-by-example generated csv files with '
             'invert_imagery_probability `{}`, min_num_examples `{}`, and '
             'example_class_selection `{}`. Please run the ./filter_examples.py '
             'script to generate the missing data, and use the `cache_dir` '
             'argument to this method to specify the location of the generated '
             'data csv files.'.format(invert_imagery_probability,
                                      min_num_examples, example_class_selection))
      raise ValueError(msg)
  
    return get_filtered_client_data_for_training(
        path_to_read_inversions_csv, path_to_read_example_indices_csv, batch_size)
  
  
  def get_filtered_client_data_for_training(path_to_read_inversions_csv,
                                            path_to_read_example_indices_csv,
                                            batch_size):
    """Form ClientData using paths to pixel inversion, example selection data."""
  
    raw_client_data = emnist_data_utils.create_real_images_tff_client_data(
        'train')
    client_ids = raw_client_data.client_ids
  
    selected_client_ids_inversion_map = None
    client_ids_example_indices_map = None
    # If filter-by-user or filter-by-example, load the csv data into maps, and
    # update the client IDs to just the users that will be part of training.
    if path_to_read_inversions_csv is not None:
      selected_client_ids_inversion_map, client_ids_example_indices_map = (
          _get_client_ids_inversion_and_example_indices_maps(
              path_to_read_inversions_csv, path_to_read_example_indices_csv))
      client_ids = list(selected_client_ids_inversion_map.keys())
  
    def _get_dataset(client_id):
      """Retrieve/preprocess a tf.data.Dataset for a given client_id."""
      raw_ds = raw_client_data.create_tf_dataset_for_client(client_id)
  
      invert_imagery = False
      if selected_client_ids_inversion_map:
        invert_imagery = selected_client_ids_inversion_map[client_id]
  
      # If filter-by-example, do it here.
      if client_ids_example_indices_map:
        raw_ds = _filter_by_example(raw_ds, client_ids_example_indices_map,
                                    client_id)
  
      return emnist_data_utils.preprocess_img_dataset(
          raw_ds,
          invert_imagery=invert_imagery,
          include_label=False,
          batch_size=batch_size,
          shuffle=True,
          repeat=False)
  
    return tff.simulation.datasets.ClientData.from_clients_and_fn(
        client_ids, _get_dataset)
  
  
  def _filter_by_example(raw_ds, client_ids_example_indices_map, client_id):
    """Form a tf.data.Dataset from the examples in the map for the client_id."""
    example_indices = client_ids_example_indices_map[client_id]
    # B/c the csv stores the list as a string, we need to do some slightly
    # klugey conversion from a string to list. (We strip off the first and
    # last characters in the string, which are [ and ], and then split on
    # commas as delimiters, to recover the original list of ints.
    example_indices = [int(s) for s in example_indices[1:-1].split(',')]
  
    # Get the elements (OrderedDicts) in the raw data which are at the indices
    # indicated by the list above.
    elements = []
    index = 0
    for element in raw_ds:
      if index in example_indices:
        elements.append(element)
      index += 1
  
    # Bind the elements (via a generator fn) into a new tf.data.Dataset.
    def _generator():
      for element in elements:
        yield element
  
    return tf.data.Dataset.from_generator(_generator, raw_ds.output_types,
                                          raw_ds.output_shapes)
  
  
  def _get_client_ids_inversion_and_example_indices_maps(
      path_to_read_inversions_csv, path_to_read_example_indices_csv):
    """Return paths to csv files storing maps indicating the data to train on."""
    if path_to_read_inversions_csv is None:
      raise ValueError(
          'No path provided to the CSV file that stores map from client ids to '
          'image inversion data.')
  
    # Load (from CSV file) the specific client IDs that the GAN will train on, and
    # whether or not the images on that client are inverted.
    selected_client_ids_inversion_map = {}
    with tf.io.gfile.GFile(path_to_read_inversions_csv, 'r') as csvfile:
      csvreader = csv.reader(csvfile)
      for [key, val] in csvreader:
        selected_client_ids_inversion_map[key] = (val == 'True')
  
    # If specified (via CSV file), the specific examples on each client ID that
    # the GAN will be trained on.
    client_ids_example_indices_map = None
    if path_to_read_example_indices_csv:
      client_ids_example_indices_map = {}
      with tf.io.gfile.GFile(path_to_read_example_indices_csv, 'r') as csvfile:
        csvreader = csv.reader(csvfile)
        for [key, val] in csvreader:
          client_ids_example_indices_map[key] = val
  
      set_1 = set(client_ids_example_indices_map.keys())
      set_2 = set(selected_client_ids_inversion_map.keys())
      symmetric_diff = set_1 ^ set_2
      if symmetric_diff:
        raise ValueError(
            'The CSV files at path_to_read_inversions_csv and '
            'path_to_read_example_indices_csv contain different keys.')
  
    return selected_client_ids_inversion_map, client_ids_example_indices_map
  
apis: tensorflow.squeeze
apis: tensorflow.io.gfile.GFile tensorflow.data.Dataset.from_generator