Invio di dati diversi a clienti particolari con tff.federated_select

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza la fonte su GitHub Scarica il taccuino

Questo tutorial mostra come implementare algoritmi federati personalizzati in TFF che richiedono l'invio di dati diversi a client diversi. Potresti già avere familiarità con tff.federated_broadcast che invia un singolo valore posizionato sul server a tutti i client. Questa esercitazione si concentra sui casi in cui parti diverse di un valore basato su server vengono inviate a client diversi. Ciò può essere utile per suddividere parti di un modello tra diversi client per evitare di inviare l'intero modello a un singolo client.

Iniziamo importando sia tensorflow che 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

Invio di valori diversi in base ai dati del cliente

Considera il caso in cui abbiamo una lista posizionata sul server da cui vogliamo inviare alcuni elementi a ciascun client in base ad alcuni dati posizionati sul client. Ad esempio, un elenco di stringhe sul server e sui client un elenco di indici da scaricare separati da virgole. Possiamo implementarlo come segue:

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

Quindi possiamo simulare il nostro calcolo fornendo l'elenco di stringhe posizionato sul server e i dati di stringa per ciascun 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'>]

Invio di un elemento casuale a ciascun cliente

In alternativa, può essere utile inviare una porzione casuale dei dati del server a ciascun client. Possiamo implementarlo generando prima una chiave casuale su ciascun client e quindi seguendo un processo di selezione simile a quello utilizzato sopra:

@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

Poiché la nostra funzione broadcast_random_element non accetta alcun dato posizionato sul client, dobbiamo configurare il runtime di simulazione TFF con un numero predefinito di client da utilizzare:

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

Quindi possiamo simulare la selezione. Possiamo modificare default_num_clients sopra e l'elenco di stringhe sotto per generare risultati diversi, o semplicemente rieseguire il calcolo per generare output casuali differenti.

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