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

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

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

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

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

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

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

Прежде чем мы начнем, попробуйте запустить следующий пример «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

# Must use the Python context because it
# supports tff.sequence_* intrinsics.
executor_factory = tff.framework.local_executor_factory(
    support_sequence_ops=True)
execution_context = tff.framework.ExecutionContext(
    executor_fn=executor_factory)
tff.framework.set_default_context(execution_context)
@tff.federated_computation
def hello_world():
  return 'Hello, World!'

hello_world()
b'Hello, World!'

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

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

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

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

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

mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
[(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

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

Поэтому мы настоятельно рекомендуем писать сложную логику TF как функции Python автономных (то есть, без tff.tf_computation украшения). Таким образом , логика TensorFlow может быть разработана и протестирована с использованием ТФА лучших практик и инструментов (как нетерпеливый режим) до сериализации вычислений для 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 , где объясняется , что в то время как мы можем выразить логику вычислений ПТФ с помощью Python, под капотом ПТФ вычислений не Python. Символ BATCH_TYPE определен выше , представляет собой абстрактную спецификацию типа ФАТ. Важно отличать этот абстрактный тип TFF от конкретных типов представления Python, например, контейнеров , таких как dict или collections.namedtuple , которые могут быть использованы для представления типа TFF в теле функции Python. В отличии от Python, ПТФ имеет один абстрактный тип конструктора 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 были объединены в 2-кортеж формальных параметров; Вы можете распознать тип batch_loss как (<MODEL_TYPE,BATCH_TYPE> -> float32) .

str(batch_loss.type_signature)
'(<model=<weights=float32[784,10],bias=float32[10]>,batch=<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.3025851

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

Что происходит , когда мы вызываем batch_loss ? Тело Python из batch_loss уже прослеживается и сериализовать в вышеприведенном ячейке , где она была определена. ФАТ выступает в качестве вызывающего абонента к batch_loss во время определения вычислений, а в качестве цели вызова в момент времени batch_loss вызывается. В обеих ролях TFF служит мостом между системой абстрактных типов TFF и типами представления Python. В то время вызова, ПТФ будет принимать большинство стандартных типов контейнеров Python ( dict , list , tuple , collections.namedtuple и т.д.) в качестве конкретных представлений абстрактного ПТФА кортежей. Кроме того, хотя, как отмечалось выше, вычисления 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)
'(<initial_model=<weights=float32[784,10],bias=float32[10]>,batch=<x=float32[?,784],y=int32[?]>,learning_rate=float32> -> <weights=float32[784,10],bias=float32[10]>)'

При вызове функции Python , украшенной tff.tf_computation в теле другой такой функции, логика внутреннего вычисления TFF встраивается ( в основном, встраивается) в логике внешнего. Как было отмечено выше, если вы пишете как вычисления, вероятно , предпочтительнее , чтобы сделать внутреннюю функцию ( batch_loss в данном случае) обычный Python или tf.function , а не 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.19690023, 0.13176313, 0.10113225, 0.08273812, 0.070301384]

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

Теперь, так как 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):

  @tff.tf_computation(LOCAL_DATA_TYPE, tf.float32)
  def _insert_learning_rate_to_sequence(dataset, learning_rate):
    return dataset.map(lambda x: (x, learning_rate))

  batches_with_learning_rate = _insert_learning_rate_to_sequence(all_batches, learning_rate)

  # Mapping function to apply to each batch.
  @tff.federated_computation(MODEL_TYPE, batches_with_learning_rate.type_signature.element)
  def batch_fn(model, batch_with_lr):
    batch, lr = batch_with_lr
    return batch_train(model, batch, lr)

  return tff.sequence_reduce(batches_with_learning_rate, initial_model, batch_fn)
str(local_train.type_signature)
'(<initial_model=<weights=float32[784,10],bias=float32[10]>,learning_rate=float32,all_batches=<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 -typed элементов, начальное состояние сокращения (мы будем называть его абстрактно , как нуль) некоторого типа U , а оператор сокращения типа (<U,T> -> U) , который изменяет состояние сокращения путем обработки одного элемента. Результатом является конечное состояние редукции после обработки всех элементов в последовательном порядке. В нашем примере состояние редукции - это модель, обученная на префиксе данных, а элементы - это пакеты данных.

Во- вторых, обратите внимание , что мы снова использовали один расчет ( batch_train ) в качестве компонента в другой ( local_train ), но не напрямую. Мы не можем использовать его как оператор редукции, потому что он принимает дополнительный параметр - скорость обучения. Чтобы решить эту проблему, мы определяем встроенный федеративного вычисления batch_fn , который связывается с local_train «ами параметр learning_rate в своем теле. Для дочернего вычисления, определенного таким образом, разрешено захватывать формальный параметр своего родителя, пока дочернее вычисление не вызывается за пределами тела его родителя. Вы можете думать об этой модели в качестве эквивалента 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):

  @tff.tf_computation(MODEL_TYPE, LOCAL_DATA_TYPE)
  def _insert_model_to_sequence(model, dataset):
    return dataset.map(lambda x: (model, x))

  model_plus_data = _insert_model_to_sequence(model, all_batches)

  @tff.tf_computation(tf.float32, batch_loss.type_signature.result)
  def tff_add(accumulator, arg):
    return accumulator + arg

  return tff.sequence_reduce(
      tff.sequence_map(
          batch_loss,
          model_plus_data), 0., tff_add)
str(local_eval.type_signature)
'(<model=<weights=float32[784,10],bias=float32[10]>,all_batches=<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 для встраивания TensorFlow логики в 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.43484688

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

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.type_at_server(MODEL_TYPE)
CLIENT_DATA_TYPE = tff.type_at_clients(LOCAL_DATA_TYPE)

Со всеми введенными до сих пор определениями выражение объединенной оценки в 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 в более простых сценариев, и на интуитивном уровне, они работают , как и ожидалось, но есть еще в этом разделе кода , чем встречает глаз, так что давайте идти по ней осторожно.

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

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

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

В самом деле, это пример неявного приведения типа в 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.type_at_server(tf.float32)


@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.60552215576172
round 1, loss=20.365678787231445
round 2, loss=19.27480125427246
round 3, loss=18.311111450195312
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 как более полный пример, и как способ , чтобы продемонстрировать некоторые из практики кодирования мы хотели бы поощрять.