TFF是不追踪或者维持tff.CLIENTS端任何东西的，在之前的博客TFF Frame的Case代码里，用户的模型是在拿到server_message后重建的（也就是只训练几轮local_model_variables）而不是一直保持、追踪的。发这篇reconstrction文章的谷歌那批人说，因为现实应用的限制，所以stateless更好，然后TFF也是Stateless。很多之前的personalization的文章都是要维持tff.CLIENTS端状态的，这篇文章就简单介绍一下如果把TFF做成stateful主要参考了为TFF提供的源代码^[1]，这里选择仍然是之前的矩阵分解进行电影推荐的例子，里面的tff.learning.Model和tff.template.IterativeProcess都要重写，正好可以当作TFF Frame和TFF Core的实践。

Model

先把需要用到的深度模型写出来，直接用tf.keras.layers提供的方法，要是按照TFF Frame里从最底层搭建太麻烦了

class UserEmbedding(tf.keras.layers.Layer):
    def __init__(self, num_latent_factors, **kwargs) -> None:
        super().__init__(**kwargs)
        self.num_latent_factors = num_latent_factors
    
    def build(self, input_shape):
        self.embedding = self.add_weight(
            shape=(1, self.num_latent_factors), 
            initializer='uniform', 
            dtype=tf.float32, 
            name='UserEmbeddingKernel'
        )
        super().build(input_shape)
    
    def call(self, inputs):
        return self.embedding
    
    def compute_output_shape(self):
        return (1, self.num_latent_factors)

def tf_model(
    num_items: int, 
    num_latent_factors: int) -> tf.keras.Model:
    item_input = tf.keras.layers.Input(shape=[1], name='Item')
    item_embedding_layer = tf.keras.layers.Embedding(
        num_items, 
        num_latent_factors, 
        name="ItemEmbedding"
    )
    flat_item_vec = tf.keras.layers.Flatten(name="FlattenItem")(
        item_embedding_layer(item_input)
    )

    user_embedding_layer = UserEmbedding(
        num_latent_factors, 
        name="UserEmbedding"
    )
    flat_user_vec = user_embedding_layer(item_input)

    pred = tf.keras.layers.Dot(axes=1, normalize=False, name='Dot')([
        flat_user_vec, flat_item_vec
    ])
    model = tf.keras.Model(inputs=item_input, outputs=pred)
    return model

model_fn = functools.partial(tf_model, num_items=3706, num_latent_factors=50)

然后开始根据实际需要编写创建变量、计算预测结果、前向传播、本地评估和federated_output_computation要暴露的protocol。这里我也很懒，模型的变量统统塞给model这个参数、能调用keras接口的就绝对不运算。

ModelVariables = collections.namedtuple('ModelVariables', 'num_examples, loss_sum')
loss_fn = tf.keras.losses.MeanSquaredError()
metrics_fn = lambda: tf.keras.metrics.Mean()

def create_variables():
    return ModelVariables(
        num_examples = tf.Variable(0.0, name="num_examples", trainable=False), 
        loss_sum = tf.Variable(0.0, name="loss_sum", trainable=False)
    )

def predict_on_batch(model, x):
    return model(x)

def tf_forward_pass(model, variables, batch):
    pred = model(batch['x'])
    loss = loss_fn(batch['y'], pred)
    num_examples = tf.cast(tf.size(batch['y']), tf.float32)
    variables.num_examples.assign_add(num_examples)
    variables.loss_sum.assign_add(loss * num_examples)
    return loss, pred

def get_local_metrics(variables):
    return collections.OrderedDict(
        num_examples = variables.num_examples, 
        loss = variables.loss_sum / variables.num_examples
    )

@tff.federated_computation
def aggregate_metrics_across_clients(metrics):
    return collections.OrderedDict(
        num_examples = tff.federated_sum(metrics.num_examples), 
        loss = tff.federated_mean(metrics.loss, metrics.num_examples)
    )

这里自己编写的类比之前多添了一个_model属性，同时为了维持一些层的参数，增加了local_model_variables和global_model_variables两个属性。

class tff_model(tff.learning.Model):

    def __init__(self) -> None:
        self._variables = create_variables()
        self._model = model_fn()
        self._model.build(input_shape=tf_model_input_shape)
    
    @property
    def trainable_variables(self):
        return self._model.trainable_variables
    
    @property
    def non_trainable_variables(self):
        return self._model.non_trainable_variables

    @property
    def local_variables(self):
        return [self._variables.num_examples, self._variables.loss_sum, 
        self._variables.accuracy_sum]

    @property
    def local_model_variables(self):
        return self._model.get_layer(name="UserEmbedding").weights
    
    @property
    def global_model_variables(self):
        return self._model.get_layer(name="ItemEmbedding").weights
    
    @property
    def input_spec(self):
        return tff_model_input_spec
    
    @tf.function
    def predict_on_batch(self, x, training=True):
        del training
        return predict_on_batch(self._model, x)
    
    @tf.function
    def forward_pass(self, batch, training=True):
        del training
        loss, predictions = tf_forward_pass(self._model, self._variables, batch)
        num_examples = tf.shape(batch['x'])[0]
        return tff.learning.BatchOutput(
            loss = loss, 
            predictions = predictions, 
            num_examples = num_examples
        )
    
    @tf.function
    def report_local_outputs(self):
        return get_local_metrics(self._variables)
    
    @property
    def federated_output_computation(self):
        return aggregate_metrics_across_clients
    
    @tf.function
    def report_local_unfinalized_metrics(self):
        return collections.OrderedDict(
                num_examples=[self._variables.num_examples],
                loss=[self._variables.loss_sum, self._variables.num_examples]
                )
    
    def metric_finalizers(self):
        return collections.OrderedDict(
            num_examples=tf.function(func=lambda x: x[0]),
            loss=tf.function(func=lambda x: x[0] / x[1])
        )

'''测试代码
a = tff_model()
iterative_process = tff.learning.build_federated_averaging_process(
    tff_model,
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02))
state = iterative_process.initialize()
state, metrics = iterative_process.next(state, tf_test_datasets[:5])
'''

IterativeProcess

整个IterativeProcess的框架大致是这个样子的：

算了太困了，先睡觉了，先传个代码吧。

Container

Code

import nest_asyncio
nest_asyncio.apply()

import collections 
import functools
import io 
import os 
import requests 
import zipfile 
from typing import List, Optional, Tuple
import attr

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf 
import tensorflow_federated as tff 

# data preparation
def download_movielens_data(dataset_path):
  """Downloads and copies MovieLens data to local /tmp directory."""
  if dataset_path.startswith('http'):
    r = requests.get(dataset_path)
    z = zipfile.ZipFile(io.BytesIO(r.content))
    z.extractall(path='/tmp')
  else:
    tf.io.gfile.makedirs('/tmp/ml-1m/')
    for filename in ['ratings.dat', 'movies.dat', 'users.dat']:
      tf.io.gfile.copy(
          os.path.join(dataset_path, filename),
          os.path.join('/tmp/ml-1m/', filename),
          overwrite=True)

download_movielens_data('http://files.grouplens.org/datasets/movielens/ml-1m.zip')

def load_movielens_data(
    data_directory: str='/tmp') -> Tuple[pd.DataFrame, pd.DataFrame]:
    ratings_df = pd.read_csv(
        os.path.join(data_directory, 'ml-1m', 'ratings.dat'), 
        sep="::", 
        names=['UserID', 'MovieID', 'Rating', 'Timestamp'], 
        engine="python", 
        encoding="ISO-8859-1"
    )
    movies_df = pd.read_csv(
        os.path.join(data_directory, 'ml-1m', 'movies.dat'), 
        sep="::", 
        names=['MovieID', 'Title', 'Genres'], 
        engine='python', 
        encoding="ISO-8859-1"
    )

    movie_mapping = {
        old_movie:new_movie for new_movie, old_movie in enumerate(
            ratings_df.MovieID.astype('category').cat.categories
        )
    }
    user_mapping = {
        old_user:new_user for new_user, old_user in enumerate(
            ratings_df.UserID.astype('category').cat.categories
        )
    }
    ratings_df.MovieID = ratings_df.MovieID.map(movie_mapping)
    ratings_df.UserID = ratings_df.UserID.map(user_mapping)
    movies_df.MovieID = movies_df.MovieID.map(movie_mapping)

    movies_df = movies_df[pd.notnull(movies_df.MovieID)]
    return ratings_df, movies_df

ratings_df, movies_df = load_movielens_data()

def create_tf_datasets(ratings_df: pd.DataFrame, batch_size: int=1, max_examples_per_user: Optional[int]=None, max_clients: Optional[int]=None) -> List[tf.data.Dataset]:
    num_users = len(ratings_df)
    if max_clients is not None:
        num_users = min(num_users, max_clients)
    
    def rating_batch_map_fn(rating_batch):
        return collections.OrderedDict([
            ('x', tf.cast(rating_batch[:, 1:2], tf.int64)), 
            ('y', tf.cast(rating_batch[:, 2:3], tf.float32))
        ])
    
    tf_datasets = []
    for user_id in range(num_users):
        user_ratings_df = ratings_df[ratings_df.UserID == user_id]
        tf_dataset = tf.data.Dataset.from_tensor_slices(user_ratings_df)
        tf_dataset = tf_dataset.take(max_examples_per_user).shuffle(buffer_size=max_examples_per_user, seed=42).batch(batch_size).map(
            rating_batch_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
        )
        tf_datasets.append(tf_dataset)
    return tf_datasets

def split_tf_datasets(
    tf_datasets: List[tf.data.Dataset], 
    train_fraction: float = 0.8, 
    val_fraction: float = 0.1,
) -> Tuple[List[tf.data.Dataset], List[tf.data.Dataset], List[tf.data.Dataset]]:
    np.random.seed(42)
    np.random.shuffle(tf_datasets)

    train_idx = int(len(tf_datasets) * train_fraction)
    val_idx = int(len(tf_datasets) * (train_fraction + val_fraction))
    return (tf_datasets[:train_idx], tf_datasets[train_idx:val_idx],
            tf_datasets[val_idx:])

tf_datasets = create_tf_datasets(
    ratings_df=ratings_df,
    batch_size=5,
    max_examples_per_user=300,
    max_clients=2000)

tf_train_datasets, tf_val_datasets, tf_test_datasets = split_tf_datasets(
    tf_datasets,
    train_fraction=0.8,
    val_fraction=0.1)

tf_model_input_shape = tf_train_datasets[0].element_spec['x'].shape
tff_model_input_spec = tf_train_datasets[0].element_spec

# build model
class UserEmbedding(tf.keras.layers.Layer):
    def __init__(self, num_latent_factors, **kwargs) -> None:
        super().__init__(**kwargs)
        self.num_latent_factors = num_latent_factors
    
    def build(self, input_shape):
        self.embedding = self.add_weight(
            shape=(1, self.num_latent_factors), 
            initializer='uniform', 
            dtype=tf.float32, 
            name='UserEmbeddingKernel'
        )
        super().build(input_shape)
    
    def call(self, inputs):
        return self.embedding
    
    def compute_output_shape(self):
        return (1, self.num_latent_factors)

def tf_model(
    num_items: int, 
    num_latent_factors: int) -> tf.keras.Model:
    item_input = tf.keras.layers.Input(shape=[1], name='Item')
    item_embedding_layer = tf.keras.layers.Embedding(
        num_items, 
        num_latent_factors, 
        name="ItemEmbedding"
    )
    flat_item_vec = tf.keras.layers.Flatten(name="FlattenItem")(
        item_embedding_layer(item_input)
    )

    user_embedding_layer = UserEmbedding(
        num_latent_factors, 
        name="UserEmbedding"
    )
    flat_user_vec = user_embedding_layer(item_input)

    pred = tf.keras.layers.Dot(axes=1, normalize=False, name='Dot')([
        flat_user_vec, flat_item_vec
    ])
    model = tf.keras.Model(inputs=item_input, outputs=pred)
    return model

model_fn = functools.partial(tf_model, num_items=3706, num_latent_factors=50)

ModelVariables = collections.namedtuple('ModelVariables', 'num_examples, loss_sum')
loss_fn = tf.keras.losses.MeanSquaredError()
metrics_fn = lambda: tf.keras.metrics.Mean()

def create_variables():
    return ModelVariables(
        num_examples = tf.Variable(0.0, name="num_examples", trainable=False), 
        loss_sum = tf.Variable(0.0, name="loss_sum", trainable=False)
    )

def predict_on_batch(model, x):
    return model(x)

def tf_forward_pass(model, variables, batch):
    pred = model(batch['x'])
    loss = loss_fn(batch['y'], pred)
    num_examples = tf.cast(tf.size(batch['y']), tf.float32)
    variables.num_examples.assign_add(num_examples)
    variables.loss_sum.assign_add(loss * num_examples)
    return loss, pred

def get_local_metrics(variables):
    return collections.OrderedDict(
        num_examples = variables.num_examples, 
        loss = variables.loss_sum / variables.num_examples
    )

@tff.federated_computation
def aggregate_metrics_across_clients(metrics):
    return collections.OrderedDict(
        num_examples = tff.federated_sum(metrics.num_examples), 
        loss = tff.federated_mean(metrics.loss, metrics.num_examples)
    )

class tff_model(tff.learning.Model):

    def __init__(self) -> None:
        self._variables = create_variables()
        self._model = model_fn()
        self._model.build(input_shape=tf_model_input_shape)
    
    @property
    def trainable_variables(self):
        return self._model.trainable_variables
    
    @property
    def non_trainable_variables(self):
        return self._model.non_trainable_variables

    @property
    def local_variables(self):
        return [self._variables.num_examples, self._variables.loss_sum, 
        self._variables.accuracy_sum]

    @property
    def local_model_variables(self):
        return self._model.get_layer(name="UserEmbedding").weights
    
    @property
    def global_model_variables(self):
        return self._model.get_layer(name="ItemEmbedding").weights
    
    @property
    def input_spec(self):
        return tff_model_input_spec
    
    @tf.function
    def predict_on_batch(self, x, training=True):
        del training
        return predict_on_batch(self._model, x)
    
    @tf.function
    def forward_pass(self, batch, training=True):
        del training
        loss, predictions = tf_forward_pass(self._model, self._variables, batch)
        num_examples = tf.shape(batch['x'])[0]
        return tff.learning.BatchOutput(
            loss = loss, 
            predictions = predictions, 
            num_examples = num_examples
        )
    
    @tf.function
    def report_local_outputs(self):
        return get_local_metrics(self._variables)
    
    @property
    def federated_output_computation(self):
        return aggregate_metrics_across_clients
    
    @tf.function
    def report_local_unfinalized_metrics(self):
        return collections.OrderedDict(
                num_examples=[self._variables.num_examples],
                loss=[self._variables.loss_sum, self._variables.num_examples]
                )
    
    def metric_finalizers(self):
        return collections.OrderedDict(
            num_examples=tf.function(func=lambda x: x[0]),
            loss=tf.function(func=lambda x: x[0] / x[1])
            )

# iterative process
@attr.s(eq=False, frozen=True, slots=True)
class ClientState():
    index = attr.ib()
    local_weights = attr.ib()

@attr.s(eq=False, frozen=True, slots=True)
class ClientOutput():
    weights_delta = attr.ib()
    client_weight = attr.ib()
    loss = attr.ib()
    client_state = attr.ib()

@attr.s(eq=False, frozen=True, slots=True)
class ServerState():
    model_weights = attr.ib()

@attr.s(eq=False, frozen=True, slots=True)
class BroadcastMessage(object):
    model_weights = attr.ib()
   

@tf.function
def server_update(model, server_state, server_optimizer, weights_delta):
    weights = model.global_model_variables
    tf.nest.map_structure(lambda x, y: x.assign(y), weights, server_state.model_weights)

    neg_weights_delta = [-1.0 * x for x in weights_delta]
    server_optimizer.apply_gradients(zip(neg_weights_delta, weights))
    return tff.structure.update_struct(
        server_state, 
        model_weights = weights
    )

@tf.function
def build_server_broadcast_message(server_state):
  return BroadcastMessage(
      model_weights=server_state.model_weights)


@tf.function
def client_update(model, dataset, client_state, server_message, client_optimizer, local_optimizer):
    weights_local = model.local_model_variables
    weights_global = model.global_model_variables
    tf.nest.map_structure(
        lambda x, y: x.assign(y), 
        weights_local, 
        client_state.local_weights
    )
    tf.nest.map_structure(
        lambda x, y: x.assign(y),
        weights_global, 
        server_message.model_weights
    )
    num_examples = tf.constant(0, dtype=tf.int32)
    loss_sum = tf.constant(0.0, dtype=tf.float32)

    for batch in dataset:
        with tf.GradientTape(persistent=True) as tape:
            outputs = model.forward_pass(batch)
        grads_local = tape.gradient(outputs.loss, weights_local)
        grads_global = tape.gradient(outputs.loss, weights_global)
        
        localgrads_and_vars = zip(grads_local, weights_local)
        globalgrads_and_vars = zip(grads_global, weights_global)

        client_optimizer.apply_gradients(globalgrads_and_vars)
        local_optimizer.apply_gradients(localgrads_and_vars)
        batch_size = (tf.shape(batch['x'])[0])
        num_examples += batch_size
        loss_sum += outputs.loss * tf.cast(batch_size, tf.float32)

    weights_delta = tf.nest.map_structure(
        lambda x, y: x-y, weights_global, server_message.model_weights
    )
    client_weights = tf.cast(num_examples, tf.float32) 
    return ClientOutput(
        weights_delta, 
        client_weights,
        loss_sum / tf.cast(client_weights, dtype=tf.float32),
        ClientState(
            index = client_state.index, 
            local_weights = weights_local
        )
    )


def build_federated_averaging_process(
    model_fn, client_state_fn, 
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0), 
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.1), 
    local_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.5)):

    whimsy_model = model_fn()

    @tff.tf_computation
    def server_init():
        model = model_fn()
        return ServerState(
            model_weights = model.global_model_variables)
    
    server_state_type = server_init.type_signature.result
    model_weights_type = server_state_type.model_weights

    client_state_type = tff.framework.type_from_tensors(client_state_fn())

    @tff.tf_computation(server_state_type, model_weights_type)  # pytype: disable=attribute-error  # gen-stub-imports
    def server_update_fn(server_state, model_delta):
        model = model_fn()
        server_optimizer = server_optimizer_fn()
        return server_update(model, server_state, server_optimizer, model_delta)
    
    @tff.tf_computation(server_state_type)
    def server_message_fn(server_state):
        return build_server_broadcast_message(server_state)
    
    server_message_type = server_message_fn.type_signature.result
    tf_dataset_type = tff.SequenceType(whimsy_model.input_spec)

    @tff.tf_computation(tf_dataset_type, client_state_type, server_message_type)
    def client_update_fn(dataset, client_state, server_message):
        model = model_fn()
        client_optimizer = client_optimizer_fn()
        local_optimizer = local_optimizer_fn()
        return client_update(model, dataset, client_state, server_message, client_optimizer, local_optimizer)
    
    federated_server_state_type = tff.type_at_server(server_state_type)
    federated_dataset_type = tff.type_at_clients(tf_dataset_type)
    federated_client_state_type = tff.type_at_clients(client_state_type)

    @tff.federated_computation(federated_server_state_type, federated_dataset_type, federated_client_state_type)
    def run_one_round(server_state, federated_dataset, client_states):
        server_message = tff.federated_map(server_message_fn, server_state)
        server_message_at_client = tff.federated_broadcast(server_message)

        client_outputs = tff.federated_map(client_update_fn, (federated_dataset, client_states, server_message_at_client))
        weight_denom = client_outputs.client_weight
        round_model_delta = tff.federated_mean(client_outputs.weights_delta, weight=weight_denom)
        round_loss = tff.federated_mean(client_outputs.loss, weight=weight_denom)
        server_state = tff.federated_map(server_update_fn, (server_state, round_model_delta))

        return server_state, round_loss, client_outputs.client_state
    
    @tff.federated_computation
    def server_init_tff():
        return tff.federated_value(server_init(), tff.SERVER)
    
    return tff.templates.IterativeProcess(initialize_fn=server_init_tff, next_fn=run_one_round)    
def generate_client_state():
    model = tff_model()
    weights = model.local_model_variables
    return ClientState(index=-1, local_weights=weights)

# main
iterative_process  = build_federated_averaging_process(tff_model, generate_client_state)

server_state = iterative_process.initialize()

client_states = {i: ClientState(i, generate_client_state().local_weights) for i in range(5)}
for round_num in range(3):
    selected_dataset = tf_test_datasets[:5]
    sampled_client_states = [client_states[i] for i in range(5)]
    server_state, trained_loss, updated_client_states = iterative_process.next(
        server_state, selected_dataset, sampled_client_states
    )
    print(f'Round {round_num} training loss: {trained_loss}')
    for client_state in updated_client_states:
        client_index = client_state.index
        tf.nest.map_structure(lambda x,y: x.assign(y), 
        client_states[client_index].local_weights, client_state.local_weights)

Reference

1. Source code of stateful clients in tff repository