Эта страница была переведа с помощью Cloud Translation API.
Switch to English

Пользовательские федеративные алгоритмы, часть 2: Реализация федеративного усреднения

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub

Это руководство является второй частью серии из двух частей, в которых демонстрируется, как реализовать настраиваемые типы федеративных алгоритмов в TFF с использованием Федеративного ядра (FC) , которое служит основой для уровня федеративного обучения (FL) ( tff.learning ). .

Мы рекомендуем вам сначала прочитать первую часть этой серии , в которой представлены некоторые ключевые концепции и абстракции программирования, используемые здесь.

Эта вторая часть серии использует механизмы, представленные в первой части, для реализации простой версии федеративных алгоритмов обучения и оценки.

Мы рекомендуем вам ознакомиться с руководствами по классификации изображений и генерации текста, чтобы получить более подробное и мягкое введение в API федеративного обучения TFF, поскольку они помогут вам поместить концепции, которые мы здесь описываем, в контекст.

Прежде чем мы начнем

Прежде чем мы начнем, попробуйте запустить следующий пример «Hello World», чтобы убедиться, что ваша среда правильно настроена. Если это не сработает, обратитесь к руководству по установке за инструкциями.


!pip install --quiet --upgrade tensorflow_federated_nightly
!pip install --quiet --upgrade nest_asyncio

import nest_asyncio
nest_asyncio.apply()
import collections

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

# TODO(b/148678573,b/148685415): must use the reference context because it
# supports unbounded references and tff.sequence_* intrinsics.
tff.backends.reference.set_reference_context()
@tff.federated_computation
def hello_world():
  return 'Hello, World!'

hello_world()
'Hello, World!'

Внедрение федеративного усреднения

Как и в Федеративном обучении для классификации изображений , мы собираемся использовать пример MNIST, но, поскольку он предназначен для низкоуровневого руководства, мы собираемся обойти tff.simulation API и tff.simulation , написать исходный код модели и построить набор объединенных данных с нуля.

Подготовка наборов объединенных данных

Для демонстрации мы собираемся смоделировать сценарий, в котором у нас есть данные от 10 пользователей, и каждый из пользователей вносит свой вклад в то, как распознавать разные цифры. Это примерно так же безумно, как и получается.

Сначала загрузим стандартные данные MNIST:

mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step

[(x.dtype, x.shape) for x in mnist_train]
[(dtype('uint8'), (60000, 28, 28)), (dtype('uint8'), (60000,))]

Данные поступают в виде массивов Numpy, один с изображениями, а другой с цифровыми метками, причем первое измерение распространяется на отдельные примеры. Давайте напишем вспомогательную функцию, которая форматирует ее способом, совместимым с тем, как мы вводим объединенные последовательности в вычисления TFF, то есть в виде списка списков - внешний список охватывает пользователей (цифры), а внутренние - пакеты данных в последовательность каждого клиента. Как обычно, мы будем структурировать каждый пакет как пару тензоров с именами x и y , каждый из которых будет иметь ведущее измерение пакета. При этом мы также сведем каждое изображение в вектор из 784 элементов и масштабируем пиксели в нем до диапазона 0..1 , чтобы нам не пришлось загромождать логику модели преобразованиями данных.

NUM_EXAMPLES_PER_USER = 1000
BATCH_SIZE = 100


def get_data_for_digit(source, digit):
  output_sequence = []
  all_samples = [i for i, d in enumerate(source[1]) if d == digit]
  for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
    batch_samples = all_samples[i:i + BATCH_SIZE]
    output_sequence.append({
        'x':
            np.array([source[0][i].flatten() / 255.0 for i in batch_samples],
                     dtype=np.float32),
        'y':
            np.array([source[1][i] for i in batch_samples], dtype=np.int32)
    })
  return output_sequence


federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]

federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]

В качестве быстрой проверки рассмотрим тензор Y в последнем пакете данных, предоставленном пятым клиентом (тот, который соответствует цифре 5 ).

federated_train_data[5][-1]['y']
array([5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], dtype=int32)

Чтобы быть уверенным, давайте также посмотрим на изображение, соответствующее последнему элементу этого пакета.

from matplotlib import pyplot as plt

plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')
plt.grid(False)
plt.show()

PNG

О сочетании TensorFlow и TFF

В этом руководстве для компактности мы сразу же украшаем функции, которые вводят логику tff.tf_computation . Однако для более сложной логики это не тот шаблон, который мы рекомендуем. Отладка TensorFlow уже может быть проблемой, а отладка TensorFlow после того, как он был полностью сериализован и затем повторно импортирован, обязательно теряет некоторые метаданные и ограничивает интерактивность, что делает отладку еще более сложной задачей.

Поэтому мы настоятельно рекомендуем писать сложную логику TF как автономные функции Python (то есть без оформления tff.tf_computation ). Таким образом, логика TensorFlow может быть разработана и протестирована с использованием передовых методов и инструментов TF (например, tff.tf_computation режима) перед сериализацией вычислений для TFF (например, путем вызова tff.tf_computation с функцией Python в качестве аргумента).

Определение функции потерь

Теперь, когда у нас есть данные, давайте определим функцию потерь, которую мы можем использовать для обучения. Во-первых, давайте определим тип ввода как именованный кортеж TFF. Поскольку размер пакетов данных может различаться, мы устанавливаем размер пакета на None чтобы указать, что размер этого измерения неизвестен.

BATCH_SPEC = collections.OrderedDict(
    x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),
    y=tf.TensorSpec(shape=[None], dtype=tf.int32))
BATCH_TYPE = tff.to_type(BATCH_SPEC)

str(BATCH_TYPE)
'<x=float32[?,784],y=int32[?]>'

Вам может быть интересно, почему мы не можем просто определить обычный тип Python. Вспомните обсуждение в части 1 , где мы объяснили, что, хотя мы можем выразить логику вычислений TFF с помощью Python, внутренние вычисления TFF не являются Python. Определенный выше символ BATCH_TYPE представляет спецификацию абстрактного типа TFF. Важно отличать этот абстрактный тип TFF от конкретных типов представления Python, например, контейнеров, таких как dict или collections.namedtuple dict которые могут использоваться для представления типа TFF в теле функции Python. В отличие от Python, TFF имеет единственный конструктор абстрактного типа tff.StructType для контейнеров, подобных кортежу, с элементами, которые могут иметь индивидуальные имена или оставаться безымянными. Этот тип также используется для моделирования формальных параметров вычислений, поскольку вычисления TFF могут формально объявлять только один параметр и один результат - вы вскоре увидите примеры этого.

Теперь давайте определим тип параметров модели TFF, снова как именованный кортеж весов и смещения TFF.

MODEL_SPEC = collections.OrderedDict(
    weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),
    bias=tf.TensorSpec(shape=[10], dtype=tf.float32))
MODEL_TYPE = tff.to_type(MODEL_SPEC)

print(MODEL_TYPE)
<weights=float32[784,10],bias=float32[10]>

Имея эти определения, теперь мы можем определить потери для данной модели по одной партии. Обратите внимание на использование декоратора @tf.function внутри декоратора @tff.tf_computation . Это позволяет нам писать TF с использованием семантики Python, даже если они находятся внутри контекста tf.Graph созданного декоратором tff.tf_computation .

# NOTE: `forward_pass` is defined separately from `batch_loss` so that it can 
# be later called from within another tf.function. Necessary because a
# @tf.function  decorated method cannot invoke a @tff.tf_computation.

@tf.function
def forward_pass(model, batch):
  predicted_y = tf.nn.softmax(
      tf.matmul(batch['x'], model['weights']) + model['bias'])
  return -tf.reduce_mean(
      tf.reduce_sum(
          tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))

@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
def batch_loss(model, batch):
  return forward_pass(model, batch)

Как и ожидалось, вычисление batch_loss возвращает потерю float32 данной модели и одного пакета данных. Обратите внимание, как MODEL_TYPE и BATCH_TYPE были объединены в кортеж из двух формальных параметров; вы можете распознать тип batch_loss как (<MODEL_TYPE,BATCH_TYPE> -> float32) .

str(batch_loss.type_signature)
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>> -> float32)'

В качестве проверки работоспособности давайте построим начальную модель, заполненную нулями, и вычислим потери для пакета данных, который мы визуализировали выше.

initial_model = collections.OrderedDict(
    weights=np.zeros([784, 10], dtype=np.float32),
    bias=np.zeros([10], dtype=np.float32))

sample_batch = federated_train_data[5][-1]

batch_loss(initial_model, sample_batch)
2.3025854

Обратите внимание, что мы загружаем вычисление TFF исходной моделью, определенной как dict , даже несмотря на то, что тело функции Python, которая ее определяет, потребляет параметры model['weight'] как model['weight'] и model['bias'] . Аргументы вызова batch_loss не просто передаются в тело этой функции.

Что происходит, когда мы вызываем batch_loss ? Тело Python для batch_loss уже было отслежено и сериализовано в указанной выше ячейке, где оно было определено. TFF действует как вызывающий batch_loss во время определения вычисления и как цель вызова во время batch_loss . В обеих ролях TFF служит мостом между системой абстрактных типов TFF и типами представления Python. Во время вызова TFF будет принимать большинство стандартных типов контейнеров Python ( dict , list , tuple , collections.namedtuple tuple и т. Д.) Как конкретные представления абстрактных кортежей TFF. Кроме того, хотя, как отмечалось выше, вычисления TFF формально принимают только один параметр, вы можете использовать знакомый синтаксис вызова Python с позиционными и / или ключевыми аргументами в случае, когда тип параметра является кортежем - он работает, как ожидалось.

Градиентный спуск на одной партии

Теперь давайте определим вычисление, которое использует эту функцию потерь для выполнения одного шага градиентного спуска. Обратите внимание, как при определении этой функции мы используем batch_loss в качестве подкомпонента. Вы можете вызвать вычисление, построенное с помощью tff.tf_computation внутри тела другого вычисления, хотя обычно в этом нет необходимости - как отмечалось выше, поскольку сериализация теряет некоторую отладочную информацию, часто для более сложных вычислений предпочтительнее записывать и тестировать все TensorFlow без декоратора tff.tf_computation .

@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
  # Define a group of model variables and set them to `initial_model`. Must
  # be defined outside the @tf.function.
  model_vars = collections.OrderedDict([
      (name, tf.Variable(name=name, initial_value=value))
      for name, value in initial_model.items()
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate)

  @tf.function
  def _train_on_batch(model_vars, batch):
    # Perform one step of gradient descent using loss from `batch_loss`.
    with tf.GradientTape() as tape:
      loss = forward_pass(model_vars, batch)
    grads = tape.gradient(loss, model_vars)
    optimizer.apply_gradients(
        zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
    return model_vars

  return _train_on_batch(model_vars, batch)
str(batch_train.type_signature)
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>,float32> -> <weights=float32[784,10],bias=float32[10]>)'

Когда вы вызываете функцию Python, украшенную tff.tf_computation в теле другой такой функции, логика вычисления внутреннего TFF встроена (по сути, встроена) в логику внешнего. Как отмечалось выше, если вы пишете оба вычисления, вероятно, предпочтительнее сделать внутреннюю функцию (в данном случае batch_loss ) обычной Python или tf.function а не tff.tf_computation . Однако здесь мы показываем, что вызов одного tff.tf_computation внутри другого в основном работает tff.tf_computation . Это может быть необходимо, если, например, у вас нет кода Python, определяющего batch_loss , а есть только его сериализованное представление TFF.

Теперь давайте применим эту функцию несколько раз к исходной модели, чтобы увидеть, уменьшаются ли потери.

model = initial_model
losses = []
for _ in range(5):
  model = batch_train(model, sample_batch, 0.1)
  losses.append(batch_loss(model, sample_batch))
losses
[0.19690022, 0.13176313, 0.10113226, 0.082738124, 0.0703014]

Градиентный спуск по последовательности локальных данных

Теперь, поскольку кажется, что batch_train работает, давайте напишем аналогичную обучающую функцию local_train которая потребляет всю последовательность всех пакетов от одного пользователя, а не только один пакет. Новое вычисление теперь должно будет использовать tff.SequenceType(BATCH_TYPE) вместо BATCH_TYPE .

LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)

@tff.federated_computation(MODEL_TYPE, tf.float32, LOCAL_DATA_TYPE)
def local_train(initial_model, learning_rate, all_batches):

  # Mapping function to apply to each batch.
  @tff.federated_computation(MODEL_TYPE, BATCH_TYPE)
  def batch_fn(model, batch):
    return batch_train(model, batch, learning_rate)

  return tff.sequence_reduce(all_batches, initial_model, batch_fn)
str(local_train.type_signature)
'(<<weights=float32[784,10],bias=float32[10]>,float32,<x=float32[?,784],y=int32[?]>*> -> <weights=float32[784,10],bias=float32[10]>)'

В этом коротком разделе кода скрыто довольно много деталей, давайте рассмотрим их одну за другой.

Во-первых, хотя мы могли бы полностью реализовать эту логику в TensorFlow, полагаясь на tf.data.Dataset.reduce для обработки последовательности аналогично тому, как мы это делали ранее, на этот раз мы решили выразить логику на связующем языке. , как tff.federated_computation . Мы использовали объединенный оператор tff.sequence_reduce для выполнения сокращения.

Оператор tff.sequence_reduce используется аналогично tf.data.Dataset.reduce . Вы можете думать об этом как о том же, что и tf.data.Dataset.reduce , но для использования внутри федеративных вычислений, которые, как вы помните, не могут содержать код TensorFlow. Это шаблонный оператор с формальным параметром 3-кортеж, который состоит из последовательности элементов типа T , начального состояния редукции (мы будем называть его абстрактно нулем ) некоторого типа U и оператора редукции тип (<U,T> -> U) который изменяет состояние редукции путем обработки одного элемента. Результатом является конечное состояние редукции после обработки всех элементов в последовательном порядке. В нашем примере состояние редукции - это модель, обученная на префиксе данных, а элементы - это пакеты данных.

Во-вторых, обратите внимание, что мы снова использовали одно вычисление ( batch_train ) как компонент в другом ( local_train ), но не напрямую. Мы не можем использовать его в качестве оператора редукции, потому что он принимает дополнительный параметр - скорость обучения. Чтобы решить эту проблему, мы определяем встроенное объединенное вычисление batch_fn которое связывается с параметром learning_rate local_train в своем теле. Для дочернего вычисления, определенного таким образом, разрешено захватывать формальный параметр своего родителя, пока дочернее вычисление не вызывается за пределами тела его родителя. Вы можете думать об этом шаблоне как об эквиваленте functools.partial в Python.

Практическое значение такого способа регистрации learning_rate , конечно же, заключается в том, что во всех пакетах используется одно и то же значение скорости обучения.

Теперь давайте попробуем недавно определенную локальную функцию обучения на всей последовательности данных от того же пользователя, который предоставил образец пакета (цифра 5 ).

locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])

Это сработало? Чтобы ответить на этот вопрос, нам нужно реализовать оценку.

Местная оценка

Вот один из способов реализации локальной оценки, суммируя потери по всем пакетам данных (мы могли бы точно так же вычислить среднее значение; мы оставим это в качестве упражнения для читателя).

@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
def local_eval(model, all_batches):
  # TODO(b/120157713): Replace with `tff.sequence_average()` once implemented.
  return tff.sequence_sum(
      tff.sequence_map(
          tff.federated_computation(lambda b: batch_loss(model, b), BATCH_TYPE),
          all_batches))
str(local_eval.type_signature)
'(<<weights=float32[784,10],bias=float32[10]>,<x=float32[?,784],y=int32[?]>*> -> float32)'

Опять же, есть несколько новых элементов, проиллюстрированных этим кодом, давайте рассмотрим их один за другим.

Во-первых, мы использовали два новых объединенных оператора для обработки последовательностей: tff.sequence_map который принимает функцию отображения T->U и последовательность T и генерирует последовательность из U полученную поточечным применением функции отображения, и tff.sequence_sum который просто добавляет все элементы. Здесь мы сопоставляем каждый пакет данных со значением потерь, а затем складываем полученные значения потерь для вычисления общих потерь.

Обратите внимание, что мы могли бы снова использовать tff.sequence_reduce , но это не лучший выбор - процесс сокращения по определению является последовательным, тогда как отображение и сумма могут вычисляться параллельно. Когда у вас есть выбор, лучше всего придерживаться операторов, которые не ограничивают выбор реализации, чтобы, когда наше вычисление TFF компилируется в будущем для развертывания в конкретной среде, можно в полной мере воспользоваться всеми потенциальными возможностями для более быстрого , более масштабируемое и более эффективное выполнение ресурсов.

Во-вторых, обратите внимание, что так же, как и в local_train , необходимая нам компонентная функция ( batch_loss ) принимает больше параметров, чем ожидает объединенный оператор ( tff.sequence_map ), поэтому мы снова определяем частичное, на этот раз встроенным, напрямую оборачивая lambda как tff.federated_computation . Использование встроенных оболочек с функцией в качестве аргумента - это рекомендуемый способ использования tff.tf_computation для встраивания логики tff.tf_computation в TFF.

Теперь посмотрим, сработало ли наше обучение.

print('initial_model loss =', local_eval(initial_model,
                                         federated_train_data[5]))
print('locally_trained_model loss =',
      local_eval(locally_trained_model, federated_train_data[5]))
initial_model loss = 23.025854
locally_trained_model loss = 0.4348469

Действительно, убыток уменьшился. Но что произойдет, если мы оценим его на данных другого пользователя?

print('initial_model loss =', local_eval(initial_model,
                                         federated_train_data[0]))
print('locally_trained_model loss =',
      local_eval(locally_trained_model, federated_train_data[0]))
initial_model loss = 23.025854
locally_trained_model loss = 74.50075

Как и ожидалось, дела пошли еще хуже. Модель была обучена распознавать 5 и никогда не видела 0 . Возникает вопрос - как обучение на местах повлияло на качество модели с глобальной точки зрения?

Федеративная оценка

Это момент в нашем путешествии, когда мы наконец возвращаемся к федеративным типам и федеративным вычислениям - теме, с которой мы начали. Вот пара определений типов TFF для модели, которая создается на сервере, и данных, которые остаются на клиентах.

SERVER_MODEL_TYPE = tff.FederatedType(MODEL_TYPE, tff.SERVER)
CLIENT_DATA_TYPE = tff.FederatedType(LOCAL_DATA_TYPE, tff.CLIENTS)

Со всеми введенными до сих пор определениями выражение объединенной оценки в TFF является однострочным - мы распространяем модель среди клиентов, позволяя каждому клиенту вызывать локальную оценку своей локальной части данных, а затем усредняем потери. Вот один из способов написать это.

@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)
def federated_eval(model, data):
  return tff.federated_mean(
      tff.federated_map(local_eval, [tff.federated_broadcast(model), data]))

Мы уже видели примеры tff.federated_mean и tff.federated_map в более простых сценариях, и на интуитивном уровне они работают, как ожидалось, но в этом разделе кода больше, чем кажется на первый взгляд, поэтому давайте рассмотрим его внимательно.

Во-первых, давайте разберемся, как разрешить каждому клиенту вызывать локальную оценку своей локальной части данных . Как вы можете вспомнить из предыдущих разделов, local_eval имеет сигнатуру типа в форме (<MODEL_TYPE, LOCAL_DATA_TYPE> -> float32) .

Объединенный оператор tff.federated_map - это шаблон, который принимает в качестве параметра кортеж из двух элементов, состоящий из функции сопоставления некоторого типа T->U и объединенного значения типа {T}@CLIENTS (т. {T}@CLIENTS С элементами-членами того же типа, что и параметр функции сопоставления), и возвращает результат типа {U}@CLIENTS .

Поскольку мы local_eval в качестве функции сопоставления, применяемой для каждого клиента, второй аргумент должен иметь федеративный тип {<MODEL_TYPE, LOCAL_DATA_TYPE>}@CLIENTS , то есть в номенклатуре предыдущих разделов он должен быть федеративным кортежем. Каждый клиент должен содержать полный набор аргументов для local_eval как составной части. Вместо этого мы скармливаем ему 2-элементный list Python. Что тут происходит?

В самом деле, это пример неявного приведения типа в TFF, подобного неявному приведению типов, с которым вы могли столкнуться в другом месте, например, когда вы передаете int в функцию, которая принимает float . На данный момент неявное приведение типов используется редко, но мы планируем сделать его более распространенным в TFF как способ минимизировать шаблонный код.

Неявное приведение, применяемое в этом случае, представляет собой эквивалентность между объединенными кортежами формы {<X,Y>}@Z и кортежами федеративных значений <{X}@Z,{Y}@Z> . Хотя формально это две сигнатуры разных типов, если смотреть на это с точки зрения программистов, каждое устройство в Z содержит две единицы данных X и Y То, что происходит здесь, мало чем отличается от zip в Python, и действительно, мы предлагаем оператор tff.federated_zip который позволяет выполнять такие преобразования явным образом. Когда tff.federated_map встречает кортеж в качестве второго аргумента, он просто вызывает для вас tff.federated_zip .

Учитывая вышеизложенное, теперь вы должны быть в состоянии распознать выражение tff.federated_broadcast(model) как представляющее значение типа TFF {MODEL_TYPE}@CLIENTS , а data как значение типа TFF {LOCAL_DATA_TYPE}@CLIENTS (или просто CLIENT_DATA_TYPE ) , оба фильтруются вместе с помощью неявного tff.federated_zip чтобы сформировать второй аргумент для tff.federated_map .

Оператор tff.federated_broadcast , как и следовало ожидать, просто передает данные с сервера клиентам.

Теперь давайте посмотрим, как наши локальные тренировки повлияли на средние потери в системе.

print('initial_model loss =', federated_eval(initial_model,
                                             federated_train_data))
print('locally_trained_model loss =',
      federated_eval(locally_trained_model, federated_train_data))
initial_model loss = 23.025852
locally_trained_model loss = 54.432625

Действительно, как и ожидалось, убыток увеличился. Чтобы улучшить модель для всех пользователей, нам нужно будет тренироваться на всех данных.

Федеративное обучение

Самый простой способ реализовать федеративное обучение - обучить на месте, а затем усреднить модели. Здесь используются те же строительные блоки и шаблоны, которые мы уже обсуждали, как вы можете видеть ниже.

SERVER_FLOAT_TYPE = tff.FederatedType(tf.float32, tff.SERVER)


@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,
                           CLIENT_DATA_TYPE)
def federated_train(model, learning_rate, data):
  return tff.federated_mean(
      tff.federated_map(local_train, [
          tff.federated_broadcast(model),
          tff.federated_broadcast(learning_rate), data
      ]))

Обратите внимание, что в полнофункциональной реализации федеративного усреднения, предоставляемой tff.learning , вместо усреднения моделей мы предпочитаем усреднять дельты моделей по ряду причин, например, возможность обрезать нормы обновления для сжатия и т. Д. .

Давайте посмотрим, работает ли тренировка, запустив несколько раундов тренировок и сравнив средние потери до и после.

model = initial_model
learning_rate = 0.1
for round_num in range(5):
  model = federated_train(model, learning_rate, federated_train_data)
  learning_rate = learning_rate * 0.9
  loss = federated_eval(model, federated_train_data)
  print('round {}, loss={}'.format(round_num, loss))
round 0, loss=21.60552406311035
round 1, loss=20.365678787231445
round 2, loss=19.27480125427246
round 3, loss=18.31110954284668
round 4, loss=17.45725440979004

Для полноты картины давайте теперь также запустим тестовые данные, чтобы убедиться, что наша модель хорошо обобщается.

print('initial_model test loss =',
      federated_eval(initial_model, federated_test_data))
print('trained_model test loss =', federated_eval(model, federated_test_data))
initial_model test loss = 22.795593
trained_model test loss = 17.278767

На этом наш урок завершен.

Конечно, наш упрощенный пример не отражает ряд вещей, которые вам нужно было бы сделать в более реалистичном сценарии - например, мы не вычисляли других показателей, кроме потерь. Мы рекомендуем вам изучить реализацию федеративного усреднения в tff.learning как более полный пример и как способ продемонстрировать некоторые методы кодирования, которые мы хотели бы поощрять.