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 em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Enviando dados diferentes para determinados clientes com tff.federated_select

Este tutorial demonstra como implementar algoritmos federados no TFF que exigem o envio de dados diferentes para clientes diferentes. Talvez você já conheça tff.federated_broadcast, que envia um único valor localizado no servidor a todos os clientes. O foco deste tutorial são os casos em que diferentes partes de um valor baseado no servidor são enviadas para clientes diferentes, o que pode ser útil para dividir partes de um modelo em diferentes clientes para evitar o envio de todo o modelo a um único cliente qualquer.

Vamos começar importando tensorflow e 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()

Envio de valores diferentes baseado nos dados dos clientes

Considere o caso em que temos uma lista localizada no servidor e que queremos enviar alguns elementos dela para cada cliente baseado em alguns dados localizados nos clientes. Por exemplo, uma lista de strings no servidor e, nos clientes, uma lista de índices separados por vírgula para baixar. Podemos implementar da seguinte forma:

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

Em seguida, podemos simular a computação fornecendo a lista de strings localizada no servidor bem como os dados de strings 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'>]

Envio de um elemento randomizado para cada cliente

Alternativamente, pode ser útil enviar uma parte aleatória dos dados do servidor para cada cliente. Podemos implementar primeiro gerando uma chave aleatória em cada cliente e depois seguindo um processo de seleção similar ao usado acima:

@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

Como a função broadcast_random_element não recebe dados localizados no cliente, precisamos configurar um número de clientes a serem usados no runtime de simulação do TFF:

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

Em seguida, podemos simular a seleção. Podemos alterar default_num_clients acima e a lista de strings abaixo para gerar resultados diferentes, ou simplesmente executar de novo a computação para gerar saídas aleatórias diferentes.

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