用於文字生成之聯邦學習

在 TensorFlow.org 上檢視

在 Google Colab 中執行

在 GitHub 上檢視原始碼

下載筆記本

本教學課程以上一個教學課程「用於圖片分類之聯邦學習」中的概念為基礎，並示範聯邦學習的其他幾種實用方法。

特別是，我們會載入先前訓練的 Keras 模型，並在 (模擬) 去中心化資料集上使用聯邦訓練加以精煉。這在實務上非常重要，原因如下。使用序列化模型的能力可讓聯邦學習與其他 ML 方法輕鬆結合。此外，這也允許使用越來越多的預先訓練模型，例如，從頭開始訓練語言模型很少是必要的，因為現在已有許多預先訓練模型廣泛可用 (請參閱 TF Hub)。相反地，從預先訓練模型開始，並使用聯邦學習加以精煉，以適應特定應用程式的去中心化資料的特定特性，會更有意義。

在本教學課程中，我們會從產生 ASCII 字元的 RNN 開始，並透過聯邦學習加以精煉。我們也會示範如何將最終權重回饋至原始 Keras 模型，以便使用標準工具輕鬆評估和產生文字。

pip install --quiet --upgrade tensorflow-federated

import collections
import functools
import os
import time

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

np.random.seed(0)

# Test that TFF is working:
tff.federated_computation(lambda: 'Hello, World!')()

b'Hello, World!'

載入預先訓練的模型

我們會載入一個模型，此模型是依照 TensorFlow 教學課程「使用 RNN 和即時執行功能產生文字」預先訓練的。但是，我們並非在 莎士比亞全集 上訓練，而是在查爾斯·狄更斯 雙城記 和 小氣財神 的文字上預先訓練模型。

除了擴充詞彙表之外，我們沒有修改原始教學課程，因此這個初始模型並非最先進的模型，但它會產生合理的預測，並且足以用於我們的教學課程目的。最終模型是使用 tf.keras.models.save_model(include_optimizer=False) 儲存的。

在本教學課程中，我們將使用聯邦學習來針對莎士比亞微調此模型，並使用 TFF 提供的資料之聯邦版本。

產生詞彙查閱表

# A fixed vocabularly of ASCII chars that occur in the works of Shakespeare and Dickens:
vocab = list('dhlptx@DHLPTX $(,048cgkoswCGKOSW[_#\'/37;?bfjnrvzBFJNRVZ"&*.26:\naeimquyAEIMQUY]!%)-159\r')

# Creating a mapping from unique characters to indices
char2idx = {u:i for i, u in enumerate(vocab)}
idx2char = np.array(vocab)

載入預先訓練的模型並產生一些文字

def load_model(batch_size):
  urls = {
      1: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch1.kerasmodel',
      8: 'https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel'}
  assert batch_size in urls, 'batch_size must be in ' + str(urls.keys())
  url = urls[batch_size]
  local_file = tf.keras.utils.get_file(os.path.basename(url), origin=url)  
  return tf.keras.models.load_model(local_file, compile=False)

def generate_text(model, start_string):
  # From https://tensorflow.dev.org.tw/tutorials/sequences/text_generation
  num_generate = 200
  input_eval = [char2idx[s] for s in start_string]
  input_eval = tf.expand_dims(input_eval, 0)
  text_generated = []
  temperature = 1.0

  model.reset_states()
  for i in range(num_generate):
    predictions = model(input_eval)
    predictions = tf.squeeze(predictions, 0)
    predictions = predictions / temperature
    predicted_id = tf.random.categorical(
        predictions, num_samples=1)[-1, 0].numpy()
    input_eval = tf.expand_dims([predicted_id], 0)
    text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))

# Text generation requires a batch_size=1 model.
keras_model_batch1 = load_model(batch_size=1)
print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))

Downloading data from https://storage.googleapis.com/tff-models-public/dickens_rnn.batch1.kerasmodel
16193984/16193984 [==============================] - 0s 0us/step
What of TensorFlow Federated, you ask? Same yee you? Have I so,
often games in a man who rode one knee over his friend, with the
stone faces of the dread prisoners, dud a tender mastery. They
are not alive is infirmed us--to ever resume

載入及預先處理聯邦莎士比亞資料

tff.simulation.datasets 套件提供各種資料集，這些資料集會分割成「用戶端」，其中每個用戶端都對應至可能參與聯邦學習之特定裝置上的資料集。

這些資料集提供實際的非 IID 資料分佈，可在模擬中重現真實去中心化資料訓練的挑戰。此資料的某些預先處理作業是使用 Leaf 專案 (github) 的工具完成的。

train_data, test_data = tff.simulation.datasets.shakespeare.load_data()

shakespeare.load_data() 提供的資料集包含字串 Tensors 序列，每個字串 Tensor 代表莎士比亞戲劇中特定角色說出的一行台詞。用戶端金鑰包含劇名與角色名稱的組合，因此例如 MUCH_ADO_ABOUT_NOTHING_OTHELLO 對應至戲劇無事生非中角色奧賽羅的台詞。請注意，在真實的聯邦學習情境中，永遠不會透過 ID 識別或追蹤用戶端，但在模擬中，使用鍵控資料集會很有用。

在這裡，例如，我們可以查看一些來自李爾王 (King Lear) 的資料

# Here the play is "The Tragedy of King Lear" and the character is "King".
raw_example_dataset = train_data.create_tf_dataset_for_client(
    'THE_TRAGEDY_OF_KING_LEAR_KING')
# To allow for future extensions, each entry x
# is an OrderedDict with a single key 'snippets' which contains the text.
for x in raw_example_dataset.take(2):
  print(x['snippets'])

tf.Tensor(b'', shape=(), dtype=string)
tf.Tensor(b'What?', shape=(), dtype=string)

我們現在使用 tf.data.Dataset 轉換來準備此資料，以用於訓練上方載入的字元 RNN。

# Input pre-processing parameters
SEQ_LENGTH = 100
BATCH_SIZE = 8
BUFFER_SIZE = 100  # For dataset shuffling

# Construct a lookup table to map string chars to indexes,
# using the vocab loaded above:
table = tf.lookup.StaticHashTable(
    tf.lookup.KeyValueTensorInitializer(
        keys=vocab, values=tf.constant(list(range(len(vocab))),
                                       dtype=tf.int64)),
    default_value=0)


def to_ids(x):
  s = tf.reshape(x['snippets'], shape=[1])
  chars = tf.strings.bytes_split(s).values
  ids = table.lookup(chars)
  return ids


def split_input_target(chunk):
  input_text = tf.map_fn(lambda x: x[:-1], chunk)
  target_text = tf.map_fn(lambda x: x[1:], chunk)
  return (input_text, target_text)


def preprocess(dataset):
  return (
      # Map ASCII chars to int64 indexes using the vocab
      dataset.map(to_ids)
      # Split into individual chars
      .unbatch()
      # Form example sequences of SEQ_LENGTH +1
      .batch(SEQ_LENGTH + 1, drop_remainder=True)
      # Shuffle and form minibatches
      .shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
      # And finally split into (input, target) tuples,
      # each of length SEQ_LENGTH.
      .map(split_input_target))

請注意，在原始序列的形成以及上方批次的形成中，為了簡化起見，我們使用 drop_remainder=True。這表示任何文字字元 (用戶端) 如果沒有至少 (SEQ_LENGTH + 1) * BATCH_SIZE 個字元的文字，就會有空的資料集。解決此問題的典型方法是使用特殊符記填補批次，然後遮罩損失，以避免將填補符記納入考量。

這會使範例變得稍微複雜，因此在本教學課程中，我們只使用完整批次，如同 標準教學課程 中所述。但是，在聯邦設定中，這個問題更為重要，因為許多使用者可能只有小型資料集。

現在我們可以預先處理 raw_example_dataset，並檢查類型

example_dataset = preprocess(raw_example_dataset)
print(example_dataset.element_spec)

(TensorSpec(shape=(8, 100), dtype=tf.int64, name=None), TensorSpec(shape=(8, 100), dtype=tf.int64, name=None))

編譯模型並在預先處理的資料上進行測試

我們載入了一個未編譯的 keras 模型，但為了執行 keras_model.evaluate，我們需要使用損失和指標來編譯它。我們也會編譯最佳化工具，此工具將在聯邦學習中用作裝置端最佳化工具。

原始教學課程沒有字元層級準確度 (最高機率放在正確下一個字元上的預測比例)。這是一個有用的指標，因此我們新增了它。但是，我們需要為此定義新的指標類別，因為我們的預測具有等級 3 (每個 BATCH_SIZE * SEQ_LENGTH 預測的 logits 向量)，而 SparseCategoricalAccuracy 僅預期等級 2 預測。

class FlattenedCategoricalAccuracy(tf.keras.metrics.SparseCategoricalAccuracy):

  def __init__(self, name='accuracy', dtype=tf.float32):
    super().__init__(name, dtype=dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_true = tf.reshape(y_true, [-1, 1])
    y_pred = tf.reshape(y_pred, [-1, len(vocab), 1])
    return super().update_state(y_true, y_pred, sample_weight)

現在我們可以編譯模型，並在我們的 example_dataset 上評估它。

BATCH_SIZE = 8  # The training and eval batch size for the rest of this tutorial.
keras_model = load_model(batch_size=BATCH_SIZE)
keras_model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[FlattenedCategoricalAccuracy()])

# Confirm that loss is much lower on Shakespeare than on random data
loss, accuracy = keras_model.evaluate(example_dataset.take(5), verbose=0)
print(
    'Evaluating on an example Shakespeare character: {a:3f}'.format(a=accuracy))

# As a sanity check, we can construct some completely random data, where we expect
# the accuracy to be essentially random:
random_guessed_accuracy = 1.0 / len(vocab)
print('Expected accuracy for random guessing: {a:.3f}'.format(
    a=random_guessed_accuracy))
random_indexes = np.random.randint(
    low=0, high=len(vocab), size=1 * BATCH_SIZE * (SEQ_LENGTH + 1))
data = collections.OrderedDict(
    snippets=tf.constant(
        ''.join(np.array(vocab)[random_indexes]), shape=[1, 1]))
random_dataset = preprocess(tf.data.Dataset.from_tensor_slices(data))
loss, accuracy = keras_model.evaluate(random_dataset, steps=10, verbose=0)
print('Evaluating on completely random data: {a:.3f}'.format(a=accuracy))

Downloading data from https://storage.googleapis.com/tff-models-public/dickens_rnn.batch8.kerasmodel
16193984/16193984 [==============================] - 0s 0us/step
Evaluating on an example Shakespeare character: 0.45.000
Expected accuracy for random guessing: 0.012
Evaluating on completely random data: 0.011

使用聯邦學習微調模型

TFF 會序列化所有 TensorFlow 運算，以便它們可能在非 Python 環境中執行 (即使目前只有以 Python 實作的模擬執行階段可用)。即使我們在即時模式 (TF 2.0) 中執行，目前 TFF 也會藉由在「with tf.Graph.as_default()」陳述式的內容中建構必要的運算元來序列化 TensorFlow 運算。因此，我們需要提供一個函式，TFF 可以使用此函式將我們的模型引入其控制的圖表中。我們執行方式如下

# Clone the keras_model inside `create_tff_model()`, which TFF will
# call to produce a new copy of the model inside the graph that it will 
# serialize. Note: we want to construct all the necessary objects we'll need 
# _inside_ this method.
def create_tff_model():
  # TFF uses an `input_spec` so it knows the types and shapes
  # that your model expects.
  input_spec = example_dataset.element_spec
  keras_model_clone = tf.keras.models.clone_model(keras_model)
  return tff.learning.models.from_keras_model(
      keras_model_clone,
      input_spec=input_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[FlattenedCategoricalAccuracy()])

現在我們已準備好建構聯邦平均迭代程序，我們將使用此程序來改進模型 (如需聯邦平均演算法的詳細資訊，請參閱論文「來自去中心化資料之深度網路的通訊效率學習」)。

在每輪聯邦訓練之後，我們使用編譯的 Keras 模型執行標準 (非聯邦式) 評估。當執行模擬聯邦學習且有標準測試資料集時，這對於研究目的很有用。

在實際的生產環境設定中，可能會使用相同的技術來取得以聯邦學習訓練的模型，並在集中式基準資料集上評估它們，以進行測試或品質保證。

# This command builds all the TensorFlow graphs and serializes them: 
fed_avg = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn=create_tff_model,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.5))

以下是最簡單的可能迴圈，我們在單一批次上針對單一用戶端執行一輪聯邦平均

state = fed_avg.initialize()
result = fed_avg.next(state, [example_dataset.take(5)])
state = result.state
train_metrics = result.metrics['client_work']['train']
print('loss={l:.3f}, accuracy={a:.3f}'.format(
    l=train_metrics['loss'], a=train_metrics['accuracy']))

loss=4.399, accuracy=0.139

現在讓我們編寫一個稍微有趣的訓練和評估迴圈。

為了讓此模擬仍然相對快速地執行，我們在每輪訓練中都在相同的 3 個用戶端上進行訓練，每個用戶端僅考量兩個迷你批次。

def data(client, source=train_data):
  return preprocess(source.create_tf_dataset_for_client(client)).take(5)


clients = [
    'ALL_S_WELL_THAT_ENDS_WELL_CELIA', 'MUCH_ADO_ABOUT_NOTHING_OTHELLO',
]

train_datasets = [data(client) for client in clients]

# We concatenate the test datasets for evaluation with Keras by creating a 
# Dataset of Datasets, and then identity flat mapping across all the examples.
test_dataset = tf.data.Dataset.from_tensor_slices(
    [data(client, test_data) for client in clients]).flat_map(lambda x: x)

fed_avg.initialize() 產生的模型初始狀態是根據 Keras 模型的隨機初始設定項，而不是載入的權重，因為 clone_model() 不會複製權重。若要從預先訓練的模型開始訓練，我們會直接從載入的模型在伺服器狀態中設定模型權重。

NUM_ROUNDS = 5

# The state of the FL server, containing the model and optimization state.
state = fed_avg.initialize()

# Load our pre-trained Keras model weights into the global model state.
pre_trained_weights = tff.learning.models.ModelWeights(
    trainable=[v.numpy() for v in keras_model.trainable_weights],
    non_trainable=[v.numpy() for v in keras_model.non_trainable_weights]
)
state = fed_avg.set_model_weights(state, pre_trained_weights)


def keras_evaluate(state, round_num):
  # Take our global model weights and push them back into a Keras model to
  # use its standard `.evaluate()` method.
  keras_model = load_model(batch_size=BATCH_SIZE)
  keras_model.compile(
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[FlattenedCategoricalAccuracy()])
  model_weights = fed_avg.get_model_weights(state)
  model_weights.assign_weights_to(keras_model)
  loss, accuracy = keras_model.evaluate(example_dataset, steps=2, verbose=0)
  print('\tEval: loss={l:.3f}, accuracy={a:.3f}'.format(l=loss, a=accuracy))


for round_num in range(NUM_ROUNDS):
  print('Round {r}'.format(r=round_num))
  keras_evaluate(state, round_num)
  result = fed_avg.next(state, train_datasets)
  state = result.state
  train_metrics = result.metrics['client_work']['train']
  print('\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(
      l=train_metrics['loss'], a=train_metrics['accuracy']))

print('Final evaluation')
keras_evaluate(state, NUM_ROUNDS + 1)

Round 0
    Eval: loss=3.171, accuracy=0.428
    Train: loss=4.309, accuracy=0.098
Round 1
    Eval: loss=4.188, accuracy=0.185
    Train: loss=4.037, accuracy=0.223
Round 2
    Eval: loss=3.948, accuracy=0.200
    Train: loss=3.797, accuracy=0.228
Round 3
    Eval: loss=3.826, accuracy=0.179
    Train: loss=3.662, accuracy=0.219
Round 4
    Eval: loss=3.723, accuracy=0.171
    Train: loss=3.440, accuracy=0.245
Final evaluation
    Eval: loss=3.599, accuracy=0.181

使用預設變更，我們尚未進行足夠的訓練來產生重大差異，但如果您在更多莎士比亞資料上進行更長時間的訓練，您應該會看到使用更新模型產生的文字樣式有所不同

# Set our newly trained weights back in the originally created model.
keras_model_batch1.set_weights([v.numpy() for v in keras_model.weights])
# Text generation requires batch_size=1
print(generate_text(keras_model_batch1, 'What of TensorFlow Federated, you ask? '))

What of TensorFlow Federated, you ask? She will be
heard of; or whether they recovered her faltering place, that a great mark of
being so final dark and distrustner the dearer to the chin, all
staftly towards him, or trot's in foot thro

建議的擴充功能

本教學課程只是第一步！以下是一些關於您可以嘗試擴充此筆記本的想法

• 編寫更實際的訓練迴圈，您可以在其中隨機取樣用戶端以進行訓練。
• 在用戶端資料集上使用「.repeat(NUM_EPOCHS)」，以嘗試多個本機訓練週期 (例如，如 McMahan 等人 中所述)。另請參閱「用於圖片分類之聯邦學習」，其中有說明如何執行此操作。
• 變更 compile() 命令，以實驗在用戶端上使用不同的最佳化演算法。
• 嘗試 server_optimizer 引數至 build_weighted_fed_avg，以嘗試不同的演算法來套用伺服器上的模型更新。