फ़ेडरेटेड_सेलेक्ट और विरल एकत्रीकरण के माध्यम से क्लाइंट-कुशल बड़े-मॉडल फ़ेडरेटेड लर्निंग

TensorFlow.org पर देखें Google Colab में चलाएं GitHub पर स्रोत देखें नोटबुक डाउनलोड करें

इस ट्यूटोरियल दिखाता है कि कैसे TFF एक बहुत बड़ी मॉडल जहां प्रत्येक ग्राहक के उपकरण केवल डाउनलोड प्रशिक्षित करने के लिए इस्तेमाल किया जा सकता है और मॉडल का एक छोटा सा हिस्सा अपडेट हो जाता है, का उपयोग कर tff.federated_select और विरल एकत्रीकरण। आत्म निहित, जबकि इस ट्यूटोरियल काफी है tff.federated_select ट्यूटोरियल और कस्टम FL एल्गोरिदम ट्यूटोरियल कुछ तकनीकों का यहां इस्तेमाल करने के लिए अच्छा परिचय प्रदान करते हैं।

वास्तव में, इस ट्यूटोरियल में हम मल्टी-लेबल वर्गीकरण के लिए लॉजिस्टिक रिग्रेशन पर विचार करते हैं, यह भविष्यवाणी करते हुए कि कौन से "टैग" एक बैग-ऑफ-वर्ड्स फीचर प्रतिनिधित्व के आधार पर टेक्स्ट स्ट्रिंग से जुड़े हैं। महत्वपूर्ण रूप से, संचार और क्लाइंट साइड गणना लागत एक निश्चित निरंतर (द्वारा नियंत्रित कर रहे MAX_TOKENS_SELECTED_PER_CLIENT ), और कुल मिलाकर शब्दावली आकार, जो व्यावहारिक सेटिंग में बहुत बड़ी हो सकता है के साथ पैमाने पर नहीं है।

!pip install --quiet --upgrade tensorflow-federated-nightly
!pip install --quiet --upgrade nest-asyncio

import nest_asyncio
nest_asyncio.apply()
import collections
import itertools
import numpy as np

from typing import Callable, List, Tuple

import tensorflow as tf
import tensorflow_federated as tff
tff.backends.native.set_local_python_execution_context()

प्रत्येक ग्राहक होगा federated_select ज्यादा से ज्यादा यह कई अद्वितीय टोकन मॉडल वेट की पंक्तियों। इस ग्राहक के स्थानीय मॉडल के आकार और सर्वर की राशि ऊपरी सीमा -> ग्राहक ( federated_select ) और ग्राहक -> सर्वर (federated_aggregate ) संचार का प्रदर्शन किया।

यह ट्यूटोरियल अभी भी सही ढंग से चलना चाहिए, भले ही आप इसे 1 के रूप में छोटा सेट करें (सुनिश्चित करें कि प्रत्येक क्लाइंट से सभी टोकन नहीं चुने गए हैं) या बड़े मूल्य पर, हालांकि मॉडल अभिसरण प्रभावित हो सकता है।

MAX_TOKENS_SELECTED_PER_CLIENT = 6

हम विभिन्न प्रकारों के लिए कुछ स्थिरांक भी परिभाषित करते हैं। इस colab के लिए, एक टोकन डाटासेट पार्स करने के बाद किसी विशेष शब्द के लिए एक पूर्णांक पहचानकर्ता है।

# There are some constraints on types
# here that will require some explicit type conversions:
#    - `tff.federated_select` requires int32
#    - `tf.SparseTensor` requires int64 indices.
TOKEN_DTYPE = tf.int64
SELECT_KEY_DTYPE = tf.int32

# Type for counts of token occurences.
TOKEN_COUNT_DTYPE = tf.int32

# A sparse feature vector can be thought of as a map
# from TOKEN_DTYPE to FEATURE_DTYPE. 
# Our features are {0, 1} indicators, so we could potentially
# use tf.int8 as an optimization.
FEATURE_DTYPE = tf.int32

समस्या की स्थापना: डेटासेट और मॉडल

हम इस ट्यूटोरियल में आसान प्रयोग के लिए एक छोटे से टॉय डेटासेट का निर्माण करते हैं। हालांकि, डेटासेट के प्रारूप के साथ संगत है संघीय StackOverflow , और पूर्व प्रसंस्करण और मॉडल वास्तुकला के StackOverflow टैग भविष्यवाणी समस्या से अपनाया जाता है अनुकूली संघीय अनुकूलन

डेटासेट पार्सिंग और प्री-प्रोसेसिंग

NUM_OOV_BUCKETS = 1

BatchType = collections.namedtuple('BatchType', ['tokens', 'tags'])

def build_to_ids_fn(word_vocab: List[str],
                    tag_vocab: List[str]) -> Callable[[tf.Tensor], tf.Tensor]:
  """Constructs a function mapping examples to sequences of token indices."""
  word_table_values = np.arange(len(word_vocab), dtype=np.int64)
  word_table = tf.lookup.StaticVocabularyTable(
      tf.lookup.KeyValueTensorInitializer(word_vocab, word_table_values),
      num_oov_buckets=NUM_OOV_BUCKETS)

  tag_table_values = np.arange(len(tag_vocab), dtype=np.int64)
  tag_table = tf.lookup.StaticVocabularyTable(
      tf.lookup.KeyValueTensorInitializer(tag_vocab, tag_table_values),
      num_oov_buckets=NUM_OOV_BUCKETS)

  def to_ids(example):
    """Converts a Stack Overflow example to a bag-of-words/tags format."""
    sentence = tf.strings.join([example['tokens'], example['title']],
                               separator=' ')

    # We represent that label (output tags) densely.
    raw_tags = example['tags']
    tags = tf.strings.split(raw_tags, sep='|')
    tags = tag_table.lookup(tags)
    tags, _ = tf.unique(tags)
    tags = tf.one_hot(tags, len(tag_vocab) + NUM_OOV_BUCKETS)
    tags = tf.reduce_max(tags, axis=0)

    # We represent the features as a SparseTensor of {0, 1}s.
    words = tf.strings.split(sentence)
    tokens = word_table.lookup(words)
    tokens, _ = tf.unique(tokens)
    # Note:  We could choose to use the word counts as the feature vector
    # instead of just {0, 1} values (see tf.unique_with_counts).
    tokens = tf.reshape(tokens, shape=(tf.size(tokens), 1))
    tokens_st = tf.SparseTensor(
        tokens,
        tf.ones(tf.size(tokens), dtype=FEATURE_DTYPE),
        dense_shape=(len(word_vocab) + NUM_OOV_BUCKETS,))
    tokens_st = tf.sparse.reorder(tokens_st)

    return BatchType(tokens_st, tags)

  return to_ids
def build_preprocess_fn(word_vocab, tag_vocab):

  @tf.function
  def preprocess_fn(dataset):
    to_ids = build_to_ids_fn(word_vocab, tag_vocab)
    # We *don't* shuffle in order to make this colab deterministic for
    # easier testing and reproducibility.
    # But real-world training should use `.shuffle()`.
    return dataset.map(to_ids, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  return preprocess_fn

एक छोटा खिलौना डेटासेट

हम 12 शब्दों और 3 क्लाइंट की वैश्विक शब्दावली के साथ एक छोटा खिलौना डेटासेट बनाते हैं। यह छोटा सा उदाहरण बढ़त मामलों के परीक्षण के लिए उपयोगी है (उदाहरण के लिए, हम कम से कम के साथ दो ग्राहक हैं MAX_TOKENS_SELECTED_PER_CLIENT = 6 और कोड विकासशील एक अलग टोकन, और अधिक के साथ)।

हालांकि, इस दृष्टिकोण के वास्तविक दुनिया के उपयोग के मामले 10 मिलियन या उससे अधिक की वैश्विक शब्दावली होंगे, शायद प्रत्येक क्लाइंट पर 1000 अलग-अलग टोकन दिखाई देंगे। क्योंकि डाटा का स्वरूप एक ही है, और अधिक यथार्थवादी testbed समस्याओं का विस्तार, जैसे tff.simulation.datasets.stackoverflow.load_data() डाटासेट, सरल होना चाहिए।

सबसे पहले, हम अपने शब्द और टैग शब्दसंग्रह को परिभाषित करते हैं।

# Features
FRUIT_WORDS = ['apple', 'orange', 'pear', 'kiwi']
VEGETABLE_WORDS = ['carrot', 'broccoli', 'arugula', 'peas']
FISH_WORDS = ['trout', 'tuna', 'cod', 'salmon']
WORD_VOCAB = FRUIT_WORDS + VEGETABLE_WORDS + FISH_WORDS

# Labels
TAG_VOCAB = ['FRUIT', 'VEGETABLE', 'FISH']

अब, हम छोटे स्थानीय डेटासेट के साथ 3 क्लाइंट बनाते हैं। यदि आप इस ट्यूटोरियल को कोलाब में चला रहे हैं, तो इस सेल और इसके आउटपुट को पिन करने के लिए "टैब में मिरर सेल" फीचर का उपयोग करना उपयोगी हो सकता है ताकि नीचे विकसित कार्यों के आउटपुट की व्याख्या/जांच की जा सके।

preprocess_fn = build_preprocess_fn(WORD_VOCAB, TAG_VOCAB)


def make_dataset(raw):
  d = tf.data.Dataset.from_tensor_slices(
      # Matches the StackOverflow formatting
      collections.OrderedDict(
          tokens=tf.constant([t[0] for t in raw]),
          tags=tf.constant([t[1] for t in raw]),
          title=['' for _ in raw]))
  d = preprocess_fn(d)
  return d


# 4 distinct tokens
CLIENT1_DATASET = make_dataset([
    ('apple orange apple orange', 'FRUIT'),
    ('carrot trout', 'VEGETABLE|FISH'),
    ('orange apple', 'FRUIT'),
    ('orange', 'ORANGE|CITRUS')  # 2 OOV tag
])

# 6 distinct tokens
CLIENT2_DATASET = make_dataset([
    ('pear cod', 'FRUIT|FISH'),
    ('arugula peas', 'VEGETABLE'),
    ('kiwi pear', 'FRUIT'),
    ('sturgeon', 'FISH'),  # OOV word
    ('sturgeon bass', 'FISH')  # 2 OOV words
])

# A client with all possible words & tags (13 distinct tokens).
# With MAX_TOKENS_SELECTED_PER_CLIENT = 6, we won't download the model
# slices for all tokens that occur on this client.
CLIENT3_DATASET = make_dataset([
    (' '.join(WORD_VOCAB + ['oovword']), '|'.join(TAG_VOCAB)),
    # Mathe the OOV token and 'salmon' occur in the largest number
    # of examples on this client:
    ('salmon oovword', 'FISH|OOVTAG')
])

print('Word vocab')
for i, word in enumerate(WORD_VOCAB):
  print(f'{i:2d} {word}')

print('\nTag vocab')
for i, tag in enumerate(TAG_VOCAB):
  print(f'{i:2d} {tag}')
Word vocab
 0 apple
 1 orange
 2 pear
 3 kiwi
 4 carrot
 5 broccoli
 6 arugula
 7 peas
 8 trout
 9 tuna
10 cod
11 salmon

Tag vocab
 0 FRUIT
 1 VEGETABLE
 2 FISH

इनपुट सुविधाओं (टोकन/शब्द) और लेबल (पोस्ट टैग) की कच्ची संख्या के लिए स्थिरांक परिभाषित करें। हमारी वास्तविक इनपुट / आउटपुट रिक्त स्थान हैं NUM_OOV_BUCKETS = 1 क्योंकि हम एक OOV टोकन / टैग को जोड़ने के लिए बड़ा।

NUM_WORDS = len(WORD_VOCAB) 
NUM_TAGS = len(TAG_VOCAB)

WORD_VOCAB_SIZE = NUM_WORDS + NUM_OOV_BUCKETS
TAG_VOCAB_SIZE = NUM_TAGS + NUM_OOV_BUCKETS

डेटासेट और अलग-अलग बैचों के बैच किए गए संस्करण बनाएं, जो हमारे जाते ही परीक्षण कोड में उपयोगी होंगे।

batched_dataset1 = CLIENT1_DATASET.batch(2)
batched_dataset2 = CLIENT2_DATASET.batch(3)
batched_dataset3 = CLIENT3_DATASET.batch(2)

batch1 = next(iter(batched_dataset1))
batch2 = next(iter(batched_dataset2))
batch3 = next(iter(batched_dataset3))

विरल इनपुट वाले मॉडल को परिभाषित करें

हम प्रत्येक टैग के लिए एक साधारण स्वतंत्र लॉजिस्टिक रिग्रेशन मॉडल का उपयोग करते हैं।

def create_logistic_model(word_vocab_size: int, vocab_tags_size: int):

  model = tf.keras.models.Sequential([
      tf.keras.layers.InputLayer(input_shape=(word_vocab_size,), sparse=True),
      tf.keras.layers.Dense(
          vocab_tags_size,
          activation='sigmoid',
          kernel_initializer=tf.keras.initializers.zeros,
          # For simplicity, don't use a bias vector; this means the model
          # is a single tensor, and we only need sparse aggregation of
          # the per-token slices of the model. Generalizing to also handle
          # other model weights that are fully updated 
          # (non-dense broadcast and aggregate) would be a good exercise.
          use_bias=False),
  ])

  return model

आइए सुनिश्चित करें कि यह काम करता है, पहले भविष्यवाणियां करके:

model = create_logistic_model(WORD_VOCAB_SIZE, TAG_VOCAB_SIZE)
p = model.predict(batch1.tokens)
print(p)
[[0.5 0.5 0.5 0.5]
 [0.5 0.5 0.5 0.5]]

और कुछ सरल केंद्रीकृत प्रशिक्षण:

model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.001),
              loss=tf.keras.losses.BinaryCrossentropy())
model.train_on_batch(batch1.tokens, batch1.tags)

फ़ेडरेटेड गणना के लिए बिल्डिंग ब्लॉक्स

हम का एक सरल संस्करण को लागू करेगा संघीय औसत का मुख्य अंतर होता एल्गोरिथ्म है कि प्रत्येक डिवाइस केवल मॉडल के एक प्रासंगिक सबसेट डाउनलोड करता है, और केवल उस सबसेट के लिए अद्यतन योगदान देता है।

हम का उपयोग M के लिए आशुलिपि के रूप में MAX_TOKENS_SELECTED_PER_CLIENT । उच्च स्तर पर, प्रशिक्षण के एक दौर में ये चरण शामिल होते हैं:

  1. प्रत्येक भाग लेने वाला क्लाइंट अपने स्थानीय डेटासेट पर स्कैन करता है, इनपुट स्ट्रिंग्स को पार्स करता है और उन्हें सही टोकन (इंट इंडेक्स) में मैप करता है। यह वैश्विक (बड़े) शब्दकोश के लिए उपयोग (इस संभावित का उपयोग कर बचा जा सकता है की आवश्यकता है सुविधा हैशिंग तकनीक)। फिर हम बहुत कम गिनते हैं कि प्रत्येक टोकन कितनी बार आता है। अगर U अद्वितीय टोकन डिवाइस पर पाए जाते हैं, हम चुनें num_actual_tokens = min(U, M) प्रशिक्षित करने के लिए सबसे लगातार टोकन।

  2. ग्राहकों का उपयोग federated_select के लिए मॉडल गुणांकों को पुनः प्राप्त करने num_actual_tokens सर्वर से चयनित टोकन। प्रत्येक मॉडल टुकड़ा आकार का एक टेन्सर है (TAG_VOCAB_SIZE, ) , इसलिए कुल ग्राहक को प्रेषित डेटा का आकार ज्यादा से ज्यादा है TAG_VOCAB_SIZE * M (नीचे नोट देखें)।

  3. ग्राहकों एक मानचित्रण का निर्माण global_token -> local_token जहां स्थानीय टोकन (पूर्णांक सूचकांक) का चयन किया टोकन की सूची में वैश्विक टोकन का सूचकांक है।

  4. ग्राहकों वैश्विक मॉडल केवल ज्यादा से ज्यादा के लिए गुणांक है की एक "छोटे" संस्करण का उपयोग M श्रृंखला से, टोकन [0, num_actual_tokens)global -> local मानचित्रण चयनित मॉडल स्लाइस से इस मॉडल के घने मापदंडों प्रारंभ करने में प्रयोग किया जाता है।

  5. ग्राहकों के साथ preprocessed डेटा पर SGD का उपयोग कर अपने स्थानीय मॉडल को प्रशिक्षित global -> local मानचित्रण।

  6. ग्राहकों में अपने स्थानीय मॉडल के मापदंडों बारी IndexedSlices का उपयोग कर अद्यतन local -> global पंक्तियों सूचकांक के मानचित्रण। सर्वर इन अद्यतनों को एक विरल योग एकत्रीकरण का उपयोग करके एकत्रित करता है।

  7. सर्वर उपरोक्त एकत्रीकरण का (घना) परिणाम लेता है, इसे भाग लेने वाले ग्राहकों की संख्या से विभाजित करता है, और परिणामी औसत अपडेट को वैश्विक मॉडल पर लागू करता है।

इस भाग में हम इन चरणों का है, जो तब फाइनल में जोड़ा जाएगा के लिए बिल्डिंग ब्लॉक का निर्माण federated_computation कि कैप्चर एक प्रशिक्षण दौर से भरा तर्क।

ग्राहक टोकन की गणना करें और जो मॉडल स्लाइस का फैसला federated_select

प्रत्येक डिवाइस को यह तय करने की आवश्यकता है कि मॉडल के कौन से "स्लाइस" उसके स्थानीय प्रशिक्षण डेटासेट के लिए प्रासंगिक हैं। हमारी समस्या के लिए, हम क्लाइंट प्रशिक्षण डेटा सेट में प्रत्येक टोकन के कितने उदाहरणों की गणना करके (बहुत कम!) ऐसा करते हैं।

@tf.function
def token_count_fn(token_counts, batch):
  """Adds counts from `batch` to the running `token_counts` sum."""
  # Sum across the batch dimension.
  flat_tokens = tf.sparse.reduce_sum(
      batch.tokens, axis=0, output_is_sparse=True)
  flat_tokens = tf.cast(flat_tokens, dtype=TOKEN_COUNT_DTYPE)
  return tf.sparse.add(token_counts, flat_tokens)
# Simple tests
# Create the initial zero token counts using empty tensors.
initial_token_counts = tf.SparseTensor(
    indices=tf.zeros(shape=(0, 1), dtype=TOKEN_DTYPE),
    values=tf.zeros(shape=(0,), dtype=TOKEN_COUNT_DTYPE),
    dense_shape=(WORD_VOCAB_SIZE,))

client_token_counts = batched_dataset1.reduce(initial_token_counts,
                                              token_count_fn)
tokens = tf.reshape(client_token_counts.indices, (-1,)).numpy()
print('tokens:', tokens)
np.testing.assert_array_equal(tokens, [0, 1, 4, 8])
# The count is the number of *examples* in which the token/word
# occurs, not the total number of occurences, since we still featurize
# multiple occurences in the same example as a "1".
counts = client_token_counts.values.numpy()
print('counts:', counts)
np.testing.assert_array_equal(counts, [2, 3, 1, 1])
tokens: [0 1 4 8]
counts: [2 3 1 1]

हम करने के लिए इसी मॉडल मापदंडों का चयन करेंगे MAX_TOKENS_SELECTED_PER_CLIENT सबसे अधिक बार डिवाइस पर टोकन उत्पन्न। डिवाइस पर इस सारे टोकन की तुलना में कम पाए जाते हैं, तो हम पैड सूची के उपयोग को सक्षम करने के federated_select

ध्यान दें कि अन्य रणनीतियाँ संभवतः बेहतर हैं, उदाहरण के लिए, बेतरतीब ढंग से टोकन का चयन करना (शायद उनकी घटना की संभावना के आधार पर)। यह सुनिश्चित करेगा कि मॉडल के सभी स्लाइस (जिसके लिए क्लाइंट के पास डेटा है) को अपडेट होने की कुछ संभावना है।

@tf.function
def keys_for_client(client_dataset, max_tokens_per_client):
  """Computes a set of max_tokens_per_client keys."""
  initial_token_counts = tf.SparseTensor(
      indices=tf.zeros((0, 1), dtype=TOKEN_DTYPE),
      values=tf.zeros((0,), dtype=TOKEN_COUNT_DTYPE),
      dense_shape=(WORD_VOCAB_SIZE,))
  client_token_counts = client_dataset.reduce(initial_token_counts,
                                              token_count_fn)
  # Find the most-frequently occuring tokens
  tokens = tf.reshape(client_token_counts.indices, shape=(-1,))
  counts = client_token_counts.values
  perm = tf.argsort(counts, direction='DESCENDING')
  tokens = tf.gather(tokens, perm)
  counts = tf.gather(counts, perm)
  num_raw_tokens = tf.shape(tokens)[0]
  actual_num_tokens = tf.minimum(max_tokens_per_client, num_raw_tokens)
  selected_tokens = tokens[:actual_num_tokens]
  paddings = [[0, max_tokens_per_client - tf.shape(selected_tokens)[0]]]
  padded_tokens = tf.pad(selected_tokens, paddings=paddings)
  # Make sure the type is statically determined
  padded_tokens = tf.reshape(padded_tokens, shape=(max_tokens_per_client,))

  # We will pass these tokens as keys into `federated_select`, which
  # requires SELECT_KEY_DTYPE=tf.int32 keys.
  padded_tokens = tf.cast(padded_tokens, dtype=SELECT_KEY_DTYPE)
  return padded_tokens, actual_num_tokens
# Simple test

# Case 1: actual_num_tokens > max_tokens_per_client
selected_tokens, actual_num_tokens = keys_for_client(batched_dataset1, 3)
assert tf.size(selected_tokens) == 3
assert actual_num_tokens == 3

# Case 2: actual_num_tokens < max_tokens_per_client
selected_tokens, actual_num_tokens = keys_for_client(batched_dataset1, 10)
assert tf.size(selected_tokens) == 10
assert actual_num_tokens == 4

वैश्विक टोकन को स्थानीय टोकन से मैप करें

ऊपर चयन हमें रेंज में टोकन के एक घने सेट देता है [0, actual_num_tokens) जो हम पर डिवाइस मॉडल के लिए प्रयोग करेंगे। हालांकि, डाटासेट हम पढ़ बहुत बड़ा वैश्विक शब्दावली सीमा से टोकन है [0, WORD_VOCAB_SIZE)

इस प्रकार, हमें वैश्विक टोकन को उनके संबंधित स्थानीय टोकन से मैप करने की आवश्यकता है। स्थानीय टोकन आईडी बस में अनुक्रमित द्वारा दिया जाता है selected_tokens टेन्सर पिछले चरण में गणना की।

@tf.function
def map_to_local_token_ids(client_data, client_keys):
  global_to_local = tf.lookup.StaticHashTable(
      # Note int32 -> int64 maps are not supported
      tf.lookup.KeyValueTensorInitializer(
          keys=tf.cast(client_keys, dtype=TOKEN_DTYPE),
          # Note we need to use tf.shape, not the static 
          # shape client_keys.shape[0]
          values=tf.range(0, limit=tf.shape(client_keys)[0],
                          dtype=TOKEN_DTYPE)),
      # We use -1 for tokens that were not selected, which can occur for clients
      # with more than MAX_TOKENS_SELECTED_PER_CLIENT distinct tokens.
      # We will simply remove these invalid indices from the batch below.
      default_value=-1)

  def to_local_ids(sparse_tokens):
    indices_t = tf.transpose(sparse_tokens.indices)
    batch_indices = indices_t[0]  # First column
    tokens = indices_t[1]  # Second column
    tokens = tf.map_fn(
        lambda global_token_id: global_to_local.lookup(global_token_id), tokens)
    # Remove tokens that aren't actually available (looked up as -1):
    available_tokens = tokens >= 0
    tokens = tokens[available_tokens]
    batch_indices = batch_indices[available_tokens]

    updated_indices = tf.transpose(
        tf.concat([[batch_indices], [tokens]], axis=0))
    st = tf.sparse.SparseTensor(
        updated_indices,
        tf.ones(tf.size(tokens), dtype=FEATURE_DTYPE),
        dense_shape=sparse_tokens.dense_shape)
    st = tf.sparse.reorder(st)
    return st

  return client_data.map(lambda b: BatchType(to_local_ids(b.tokens), b.tags))
# Simple test
client_keys, actual_num_tokens = keys_for_client(
    batched_dataset3, MAX_TOKENS_SELECTED_PER_CLIENT)
client_keys = client_keys[:actual_num_tokens]

d = map_to_local_token_ids(batched_dataset3, client_keys)
batch  = next(iter(d))
all_tokens = tf.gather(batch.tokens.indices, indices=1, axis=1)
# Confirm we have local indices in the range [0, MAX):
assert tf.math.reduce_max(all_tokens) < MAX_TOKENS_SELECTED_PER_CLIENT
assert tf.math.reduce_max(all_tokens) >= 0

प्रत्येक ग्राहक पर स्थानीय (उप) मॉडल को प्रशिक्षित करें

नोट federated_select एक के रूप में चयनित स्लाइस वापस आ जाएगी tf.data.Dataset चयन कुंजी के रूप में एक ही क्रम में। इसलिए, हम पहले ऐसे डेटासेट को लेने के लिए एक उपयोगिता फ़ंक्शन को परिभाषित करते हैं और इसे एकल घने टेंसर में परिवर्तित करते हैं जिसे क्लाइंट मॉडल के मॉडल वेट के रूप में उपयोग किया जा सकता है।

@tf.function
def slices_dataset_to_tensor(slices_dataset):
  """Convert a dataset of slices to a tensor."""
  # Use batching to gather all of the slices into a single tensor.
  d = slices_dataset.batch(MAX_TOKENS_SELECTED_PER_CLIENT,
                           drop_remainder=False)
  iter_d = iter(d)
  tensor = next(iter_d)
  # Make sure we have consumed everything
  opt = iter_d.get_next_as_optional()
  tf.Assert(tf.logical_not(opt.has_value()), data=[''], name='CHECK_EMPTY')
  return tensor
# Simple test
weights = np.random.random(
    size=(MAX_TOKENS_SELECTED_PER_CLIENT, TAG_VOCAB_SIZE)).astype(np.float32)
model_slices_as_dataset = tf.data.Dataset.from_tensor_slices(weights)
weights2 = slices_dataset_to_tensor(model_slices_as_dataset)
np.testing.assert_array_equal(weights, weights2)

अब हमारे पास एक साधारण स्थानीय प्रशिक्षण लूप को परिभाषित करने के लिए आवश्यक सभी घटक हैं जो प्रत्येक क्लाइंट पर चलेंगे।

@tf.function
def client_train_fn(model, client_optimizer,
                    model_slices_as_dataset, client_data,
                    client_keys, actual_num_tokens):

  initial_model_weights = slices_dataset_to_tensor(model_slices_as_dataset)
  assert len(model.trainable_variables) == 1
  model.trainable_variables[0].assign(initial_model_weights)

  # Only keep the "real" (unpadded) keys.
  client_keys = client_keys[:actual_num_tokens]

  client_data = map_to_local_token_ids(client_data, client_keys)

  loss_fn = tf.keras.losses.BinaryCrossentropy()
  for features, labels in client_data:
    with tf.GradientTape() as tape:
      predictions = model(features)
      loss = loss_fn(labels, predictions)
    grads = tape.gradient(loss, model.trainable_variables)
    client_optimizer.apply_gradients(zip(grads, model.trainable_variables))

  model_weights_delta = model.trainable_weights[0] - initial_model_weights
  model_weights_delta = tf.slice(model_weights_delta, begin=[0, 0], 
                           size=[actual_num_tokens, -1])
  return client_keys, model_weights_delta
# Simple test
# Note if you execute this cell a second time, you need to also re-execute
# the preceeding cell to avoid "tf.function-decorated function tried to 
# create variables on non-first call" errors.
on_device_model = create_logistic_model(MAX_TOKENS_SELECTED_PER_CLIENT,
                                        TAG_VOCAB_SIZE)
client_optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)
client_keys, actual_num_tokens = keys_for_client(
    batched_dataset2, MAX_TOKENS_SELECTED_PER_CLIENT)

model_slices_as_dataset = tf.data.Dataset.from_tensor_slices(
    np.zeros((MAX_TOKENS_SELECTED_PER_CLIENT, TAG_VOCAB_SIZE),
             dtype=np.float32))

keys, delta = client_train_fn(
    on_device_model,
    client_optimizer,
    model_slices_as_dataset,
    client_data=batched_dataset3,
    client_keys=client_keys,
    actual_num_tokens=actual_num_tokens)

print(delta)

कुल अनुक्रमित स्लाइस

हम का उपयोग tff.federated_aggregate के लिए एक फ़ेडरेटेड विरल राशि के निर्माण के लिए IndexedSlices । यह सरल कार्यान्वयन बाधा है कि है dense_shape पहले से स्थिर जाना जाता है। यह भी ध्यान रखें कि यह राशि केवल अर्द्ध विरल, अर्थ में है कि ग्राहक -> सर्वर संचार विरल है, लेकिन सर्वर में योग के एक घने प्रतिनिधित्व का कहना है accumulate और merge , और इस घने प्रतिनिधित्व आउटपुट।

def federated_indexed_slices_sum(slice_indices, slice_values, dense_shape):
  """
  Sumes IndexedSlices@CLIENTS to a dense @SERVER Tensor.

  Intermediate aggregation is performed by converting to a dense representation,
  which may not be suitable for all applications.

  Args:
    slice_indices: An IndexedSlices.indices tensor @CLIENTS.
    slice_values: An IndexedSlices.values tensor @CLIENTS.
    dense_shape: A statically known dense shape.

  Returns:
    A dense tensor placed @SERVER representing the sum of the client's
    IndexedSclies.
  """
  slices_dtype = slice_values.type_signature.member.dtype
  zero = tff.tf_computation(
      lambda: tf.zeros(dense_shape, dtype=slices_dtype))()

  @tf.function
  def accumulate_slices(dense, client_value):
    indices, slices = client_value
    # There is no built-in way to add `IndexedSlices`, but 
    # tf.convert_to_tensor is a quick way to convert to a dense representation
    # so we can add them.
    return dense + tf.convert_to_tensor(
        tf.IndexedSlices(slices, indices, dense_shape))


  return tff.federated_aggregate(
      (slice_indices, slice_values),
      zero=zero,
      accumulate=tff.tf_computation(accumulate_slices),
      merge=tff.tf_computation(lambda d1, d2: tf.add(d1, d2, name='merge')),
      report=tff.tf_computation(lambda d: d))

एक न्यूनतम का निर्माण federated_computation एक परीक्षण के रूप

dense_shape = (6, 2)
indices_type = tff.TensorType(tf.int64, (None,))
values_type = tff.TensorType(tf.float32, (None, 2))
client_slice_type = tff.type_at_clients(
    (indices_type, values_type))

@tff.federated_computation(client_slice_type)
def test_sum_indexed_slices(indices_values_at_client):
  indices, values = indices_values_at_client
  return federated_indexed_slices_sum(indices, values, dense_shape)

print(test_sum_indexed_slices.type_signature)
({<int64[?],float32[?,2]>}@CLIENTS -> float32[6,2]@SERVER)
x = tf.IndexedSlices(
    values=np.array([[2., 2.1], [0., 0.1], [1., 1.1], [5., 5.1]],
                    dtype=np.float32),
    indices=[2, 0, 1, 5],
    dense_shape=dense_shape)
y = tf.IndexedSlices(
    values=np.array([[0., 0.3], [3.1, 3.2]], dtype=np.float32),
    indices=[1, 3],
    dense_shape=dense_shape)

# Sum one.
result = test_sum_indexed_slices([(x.indices, x.values)])
np.testing.assert_array_equal(tf.convert_to_tensor(x), result)

# Sum two.
expected = [[0., 0.1], [1., 1.4], [2., 2.1], [3.1, 3.2], [0., 0.], [5., 5.1]]
result = test_sum_indexed_slices([(x.indices, x.values), (y.indices, y.values)])
np.testing.assert_array_almost_equal(expected, result)

एक में सब एक साथ डालने federated_computation

अब हम एक में घटक एक साथ बाध्य करने के लिए TFF का उपयोग करता tff.federated_computation

DENSE_MODEL_SHAPE = (WORD_VOCAB_SIZE, TAG_VOCAB_SIZE)
client_data_type = tff.SequenceType(batched_dataset1.element_spec)
model_type = tff.TensorType(tf.float32, shape=DENSE_MODEL_SHAPE)

हम 1.0 के सर्वर सीखने की दर के साथ अद्यतन को लागू करते हुए, फ़ेडरेटेड एवरेजिंग पर आधारित एक बुनियादी सर्वर प्रशिक्षण फ़ंक्शन का उपयोग करते हैं। यह महत्वपूर्ण है कि हम केवल क्लाइंट-आपूर्ति किए गए मॉडल के औसत के बजाय मॉडल में एक अपडेट (डेल्टा) लागू करें, अन्यथा यदि मॉडल के दिए गए स्लाइस को किसी भी क्लाइंट द्वारा किसी दिए गए दौर में प्रशिक्षित नहीं किया गया था, तो इसके गुणांक शून्य हो सकते हैं। बाहर।

@tff.tf_computation
def server_update(current_model_weights, update_sum, num_clients):
  average_update = update_sum / num_clients
  return current_model_weights + average_update

हमें कुछ और की जरूरत है tff.tf_computation घटक:

# Function to select slices from the model weights in federated_select:
select_fn = tff.tf_computation(
    lambda model_weights, index: tf.gather(model_weights, index))


# We need to wrap `client_train_fn` as a `tff.tf_computation`, making
# sure we do any operations that might construct `tf.Variable`s outside
# of the `tf.function` we are wrapping.
@tff.tf_computation
def client_train_fn_tff(model_slices_as_dataset, client_data, client_keys,
                        actual_num_tokens):
  # Note this is amaller than the global model, using
  # MAX_TOKENS_SELECTED_PER_CLIENT which is much smaller than WORD_VOCAB_SIZE.
  # W7e would like a model of size `actual_num_tokens`, but we
  # can't build the model dynamically, so we will slice off the padded
  # weights at the end.
  client_model = create_logistic_model(MAX_TOKENS_SELECTED_PER_CLIENT,
                                       TAG_VOCAB_SIZE)
  client_optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
  return client_train_fn(client_model, client_optimizer,
                         model_slices_as_dataset, client_data, client_keys,
                         actual_num_tokens)

@tff.tf_computation
def keys_for_client_tff(client_data):
  return keys_for_client(client_data, MAX_TOKENS_SELECTED_PER_CLIENT)

अब हम सभी टुकड़ों को एक साथ रखने के लिए तैयार हैं!

@tff.federated_computation(
    tff.type_at_server(model_type), tff.type_at_clients(client_data_type))
def sparse_model_update(server_model, client_data):
  max_tokens = tff.federated_value(MAX_TOKENS_SELECTED_PER_CLIENT, tff.SERVER)
  keys_at_clients, actual_num_tokens = tff.federated_map(
      keys_for_client_tff, client_data)

  model_slices = tff.federated_select(keys_at_clients, max_tokens, server_model,
                                      select_fn)

  update_keys, update_slices = tff.federated_map(
      client_train_fn_tff,
      (model_slices, client_data, keys_at_clients, actual_num_tokens))

  dense_update_sum = federated_indexed_slices_sum(update_keys, update_slices,
                                                  DENSE_MODEL_SHAPE)
  num_clients = tff.federated_sum(tff.federated_value(1.0, tff.CLIENTS))

  updated_server_model = tff.federated_map(
      server_update, (server_model, dense_update_sum, num_clients))

  return updated_server_model


print(sparse_model_update.type_signature)
(<server_model=float32[13,4]@SERVER,client_data={<tokens=<indices=int64[?,2],values=int32[?],dense_shape=int64[2]>,tags=float32[?,4]>*}@CLIENTS> -> float32[13,4]@SERVER)

आइए एक मॉडल को प्रशिक्षित करें!

अब जबकि हमारे पास अपना प्रशिक्षण कार्य है, आइए इसे आजमाते हैं।

server_model = create_logistic_model(WORD_VOCAB_SIZE, TAG_VOCAB_SIZE)
server_model.compile(  # Compile to make evaluation easy.
    optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.0),  # Unused
    loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=[ 
      tf.keras.metrics.Precision(name='precision'),
      tf.keras.metrics.AUC(name='auc'),
      tf.keras.metrics.Recall(top_k=2, name='recall_at_2'),
  ])

def evaluate(model, dataset, name):
  metrics = model.evaluate(dataset, verbose=0)
  metrics_str = ', '.join([f'{k}={v:.2f}' for k, v in 
                          (zip(server_model.metrics_names, metrics))])
  print(f'{name}: {metrics_str}')
print('Before training')
evaluate(server_model, batched_dataset1, 'Client 1')
evaluate(server_model, batched_dataset2, 'Client 2')
evaluate(server_model, batched_dataset3, 'Client 3')

model_weights = server_model.trainable_weights[0]

client_datasets = [batched_dataset1, batched_dataset2, batched_dataset3]
for _ in range(10):  # Run 10 rounds of FedAvg
  # We train on 1, 2, or 3 clients per round, selecting
  # randomly.
  cohort_size = np.random.randint(1, 4)
  clients = np.random.choice([0, 1, 2], cohort_size, replace=False)
  print('Training on clients', clients)
  model_weights = sparse_model_update(
      model_weights, [client_datasets[i] for i in clients])
server_model.set_weights([model_weights])

print('After training')
evaluate(server_model, batched_dataset1, 'Client 1')
evaluate(server_model, batched_dataset2, 'Client 2')
evaluate(server_model, batched_dataset3, 'Client 3')
Before training
Client 1: loss=0.69, precision=0.00, auc=0.50, recall_at_2=0.60
Client 2: loss=0.69, precision=0.00, auc=0.50, recall_at_2=0.50
Client 3: loss=0.69, precision=0.00, auc=0.50, recall_at_2=0.40
Training on clients [0 1]
Training on clients [0 2 1]
Training on clients [2 0]
Training on clients [1 0 2]
Training on clients [2]
Training on clients [2 0]
Training on clients [1 2 0]
Training on clients [0]
Training on clients [2]
Training on clients [1 2]
After training
Client 1: loss=0.67, precision=0.80, auc=0.91, recall_at_2=0.80
Client 2: loss=0.68, precision=0.67, auc=0.96, recall_at_2=1.00
Client 3: loss=0.65, precision=1.00, auc=0.93, recall_at_2=0.80