Gửi dữ liệu khác nhau cho khách hàng cụ thể với tff.federated_select

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Hướng dẫn này trình bày cách triển khai các thuật toán liên hợp tùy chỉnh trong TFF yêu cầu gửi dữ liệu khác nhau đến các máy khách khác nhau. Bạn có thể đã quen với tff.federated_broadcast mà sẽ gửi một giá trị máy chủ đặt duy nhất để tất cả khách hàng. Hướng dẫn này tập trung vào các trường hợp các phần khác nhau của giá trị dựa trên máy chủ được gửi đến các máy khách khác nhau. Điều này có thể hữu ích cho việc phân chia các phần của một mô hình cho các máy khách khác nhau để tránh gửi toàn bộ mô hình cho bất kỳ máy khách nào.

Hãy bắt đầu bằng cách nhập cả tensorflow và 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
tff.backends.native.set_local_python_execution_context()

Gửi các giá trị khác nhau dựa trên dữ liệu khách hàng

Hãy xem xét trường hợp chúng tôi có một số danh sách do máy chủ đặt mà từ đó chúng tôi muốn gửi một vài phần tử đến từng máy khách dựa trên một số dữ liệu do máy khách đặt. Ví dụ: danh sách các chuỗi trên máy chủ và trên máy khách, danh sách chỉ số được phân tách bằng dấu phẩy để tải xuống. Chúng ta có thể thực hiện điều đó như sau:

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

Sau đó, chúng tôi có thể mô phỏng tính toán của mình bằng cách cung cấp danh sách chuỗi do máy chủ đặt cũng như dữ liệu chuỗi cho mỗi máy khách:

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

Gửi một phần tử ngẫu nhiên cho mỗi khách hàng

Ngoài ra, có thể hữu ích khi gửi một phần dữ liệu máy chủ ngẫu nhiên đến từng máy khách. Chúng tôi có thể thực hiện điều đó bằng cách tạo khóa ngẫu nhiên trước tiên trên mỗi máy khách và sau đó thực hiện theo quy trình lựa chọn tương tự như quy trình đã sử dụng ở trên:

@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

Kể từ khi chúng tôi broadcast_random_element chức năng không mất trong bất kỳ dữ liệu khách hàng đặt, chúng ta phải cấu hình TFF Simulation Runtime với một số mặc định của khách hàng để sử dụng:

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

Sau đó, chúng ta có thể mô phỏng vùng chọn. Chúng tôi có thể thay đổi default_num_clients trên và danh sách các chuỗi dưới đây để tạo ra kết quả khác nhau, hoặc đơn giản là chạy lại việc tính toán để tạo ra ngẫu nhiên khác nhau.

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