federated_select.ipynb

Kernel: Python 3

Copyright 2021 The TensorFlow Authors.

#@title 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
#
# https://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.

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

Envío de datos diferentes a clientes particulares con tff.federated_select

Este tutorial demuestra cómo implementar algoritmos federados personalizados en TFF que requieren el envío de diferentes datos a diferentes clientes. Quizá ya esté familiarizado con tff.federated_broadcast, que envía un único valor colocado en el servidor a todos los clientes. Este tutorial se centra en casos en los que se envían diferentes partes de un valor basado en servidor a diferentes clientes. Esto puede resultar útil para dividir partes de un modelo entre diferentes clientes para evitar el envío del modelo completo a un solo cliente.

Para empezar, importemos tensorflow y tensorflow_federated.

#@test {"skip": true}
!pip install --quiet --upgrade tensorflow-federated

import tensorflow as tf
import tensorflow_federated as tff
tff.backends.native.set_sync_local_cpp_execution_context()

Envío de diferentes valores según los datos del cliente

Pensemos en el caso de que tengamos una lista colocada en el servidor desde la cual queremos enviar algunos elementos a cada cliente en función de algunos datos colocados por el cliente. Por ejemplo, una lista de cadenas en el servidor y, en los clientes, una lista separada por comas de índices para descargar. Podemos implementar eso de la siguiente manera:

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

Luego podemos simular nuestro cálculo proporcionando la lista de cadenas colocadas en el servidor, así como datos de cadena para cada cliente:

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'>]

Envío de un elemento aleatorio a cada cliente

Como alternativa, quizás sea útil enviar una parte aleatoria de los datos del servidor a cada cliente. Podemos implementar esto si generamos primero una clave aleatoria en cada cliente y luego seguimos un proceso de selección similar al que usamos anteriormente:

@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

Dado que nuestra función broadcast_random_element no acepta ningún dato colocado por el cliente, tenemos que configurar el tiempo de ejecución de simulación TFF con una cantidad predeterminada de clientes para usar:

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

Luego, podemos simular la selección. Podemos cambiar default_num_clients (arriba) y la lista de cadenas a continuación para generar resultados diferentes, o simplemente volver a ejecutar el cálculo para generar diferentes salidas aleatorias.

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