Envoi de données différentes à des clients particuliers avec tff.federated_select

Voir sur TensorFlow.org Exécuter dans Google Colab Voir la source sur GitHub Télécharger le cahier

Ce didacticiel montre comment implémenter des algorithmes fédérés personnalisés dans TFF qui nécessitent l'envoi de différentes données à différents clients. Vous pouvez déjà être familier avec tff.federated_broadcast qui envoie une seule valeur placée serveur à tous les clients. Ce didacticiel se concentre sur les cas où différentes parties d'une valeur basée sur le serveur sont envoyées à différents clients. Cela peut être utile pour diviser des parties d'un modèle entre différents clients afin d'éviter d'envoyer le modèle entier à un seul client.

Commençons par l' importation à la fois tensorflow et tensorflow_federated .

!pip install --quiet --upgrade tensorflow-federated-nightly
!pip install --quiet --upgrade nest-asyncio

import nest_asyncio
nest_asyncio.apply()
import tensorflow as tf
import tensorflow_federated as tff

Envoi de différentes valeurs en fonction des données client

Considérons le cas où nous avons une liste placée par le serveur à partir de laquelle nous voulons envoyer quelques éléments à chaque client en fonction de certaines données placées par le client. Par exemple, une liste de chaînes sur le serveur et sur les clients, une liste d'index séparés par des virgules à télécharger. Nous pouvons implémenter cela comme suit :

list_of_strings_type = tff.TensorType(tf.string, [None])
# We only ever send exactly two values to each client. The number of keys per
# client must be a fixed number across all clients.
number_of_keys_per_client = 2
keys_type = tff.TensorType(tf.int32, [number_of_keys_per_client])
get_size = tff.tf_computation(lambda x: tf.size(x))
select_fn = tff.tf_computation(lambda val, index: tf.gather(val, index))
client_data_type = tf.string

# A function from our client data to the indices of the values we'd like to
# select from the server.
@tff.tf_computation(client_data_type)
@tff.check_returns_type(keys_type)
def keys_for_client(client_string):
  # We assume our client data is a single string consisting of exactly three
  # comma-separated integers indicating which values to grab from the server.
  split = tf.strings.split([client_string], sep=',')[0]
  return tf.strings.to_number([split[0], split[1]], tf.int32)

@tff.tf_computation(tff.SequenceType(tf.string))
@tff.check_returns_type(tf.string)
def concatenate(values):
  def reduce_fn(acc, item):
    return tf.cond(tf.math.equal(acc, ''),
                   lambda: item,
                   lambda: tf.strings.join([acc, item], ','))
  return values.reduce('', reduce_fn)

@tff.federated_computation(tff.type_at_server(list_of_strings_type), tff.type_at_clients(client_data_type))
def broadcast_based_on_client_data(list_of_strings_at_server, client_data):
  keys_at_clients = tff.federated_map(keys_for_client, client_data)
  max_key = tff.federated_map(get_size, list_of_strings_at_server)
  values_at_clients = tff.federated_select(keys_at_clients, max_key, list_of_strings_at_server, select_fn)
  value_at_clients = tff.federated_map(concatenate, values_at_clients)
  return value_at_clients

Ensuite, nous pouvons simuler notre calcul en fournissant la liste de chaînes placée par le serveur ainsi que les données de chaîne pour chaque client :

client_data = ['0,1', '1,2', '2,0']
broadcast_based_on_client_data(['a', 'b', 'c'], client_data)
[<tf.Tensor: shape=(), dtype=string, numpy=b'a,b'>,
 <tf.Tensor: shape=(), dtype=string, numpy=b'b,c'>,
 <tf.Tensor: shape=(), dtype=string, numpy=b'c,a'>]

Envoi d'un élément aléatoire à chaque client

Alternativement, il peut être utile d'envoyer une partie aléatoire des données du serveur à chaque client. Nous pouvons implémenter cela en générant d'abord une clé aléatoire sur chaque client, puis en suivant un processus de sélection similaire à celui utilisé ci-dessus :

@tff.tf_computation(tf.int32)
@tff.check_returns_type(tff.TensorType(tf.int32, [1]))
def get_random_key(max_key):
  return tf.random.uniform(shape=[1], minval=0, maxval=max_key, dtype=tf.int32)

list_of_strings_type = tff.TensorType(tf.string, [None])
get_size = tff.tf_computation(lambda x: tf.size(x))
select_fn = tff.tf_computation(lambda val, index: tf.gather(val, index))

@tff.tf_computation(tff.SequenceType(tf.string))
@tff.check_returns_type(tf.string)
def get_last_element(sequence):
  return sequence.reduce('', lambda _initial_state, val: val)

@tff.federated_computation(tff.type_at_server(list_of_strings_type))
def broadcast_random_element(list_of_strings_at_server):
  max_key_at_server = tff.federated_map(get_size, list_of_strings_at_server)
  max_key_at_clients = tff.federated_broadcast(max_key_at_server)
  key_at_clients = tff.federated_map(get_random_key, max_key_at_clients)
  random_string_sequence_at_clients = tff.federated_select(
      key_at_clients, max_key_at_server, list_of_strings_at_server, select_fn)
  # Even though we only passed in a single key, `federated_select` returns a
  # sequence for each client. We only care about the last (and only) element.
  random_string_at_clients = tff.federated_map(get_last_element, random_string_sequence_at_clients)
  return random_string_at_clients

Depuis notre broadcast_random_element fonction ne prend pas toutes les données placées sur le client, nous devons configurer la TFF Simulation d' exécution avec un certain nombre de clients par défaut à utiliser:

tff.backends.native.set_local_python_execution_context(default_num_clients=3)

Ensuite, nous pouvons simuler la sélection. Nous pouvons changer default_num_clients ci - dessus et la liste des chaînes ci - dessous pour générer des résultats différents, ou simplement relancer le calcul pour générer différentes sorties aléatoires.

broadcast_random_element(tf.convert_to_tensor(['foo', 'bar', 'baz']))