TensorFlow Federated模型权重欧氏距离计算方法及联邦场景咨询

要计算TensorFlow Federated（TFF）中两个训练后模型的欧氏距离，核心是把模型权重转换为可计算的一维向量，再计算L2范数（即欧氏距离）。具体步骤如下：

1. 提取模型权重：TFF的模型权重通常以tff.learning.ModelWeights对象存在，包含trainable（可训练参数）和non_trainable（不可训练参数，如BatchNorm的均值/方差）两部分。
2. 扁平化权重：将所有参数张量展开为一维，再拼接成单个大向量。
3. 计算欧氏距离：对两个扁平化向量求差，再计算其L2范数。

代码实现 #

import tensorflow as tf
import tensorflow_federated as tff

def flatten_model_weights(weights):
    # 遍历所有可训练和不可训练参数，扁平化后拼接
    flat_tensors = []
    for param in weights.trainable:
        flat_tensors.append(tf.reshape(param, shape=[-1]))
    for param in weights.non_trainable:
        flat_tensors.append(tf.reshape(param, shape=[-1]))
    return tf.concat(flat_tensors, axis=0)

# 假设model_weights_1和model_weights_2是两个训练后的TFF模型权重对象
flat_w1 = flatten_model_weights(model_weights_1)
flat_w2 = flatten_model_weights(model_weights_2)

# 计算欧氏距离
euclidean_distance = tf.norm(flat_w1 - flat_w2).numpy()
print(f"模型间欧氏距离: {euclidean_distance:.4f}")

搭建联邦学习下的下一词预测模型，并在服务器端每轮计算接收的客户端模型权重距离，需要完成模型定义、联邦训练流程封装、权重收集与距离计算三个核心部分。

1. 定义下一词预测模型 #

以LSTM为例，用Keras搭建基础模型，再包装为TFF兼容的模型：

import tensorflow as tf
import tensorflow_federated as tff
from tensorflow.keras import Sequential, layers

def build_next_word_model(vocab_size, embedding_dim=128, rnn_units=256, seq_len=10):
    model = Sequential([
        layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim, input_length=seq_len),
        layers.LSTM(rnn_units),
        layers.Dense(vocab_size, activation='softmax')
    ])
    return model

# 包装为TFF模型
vocab_size = 10000  # 根据实际词汇表大小调整

def tff_model_fn():
    keras_model = build_next_word_model(vocab_size)
    return tff.learning.from_keras_model(
        keras_model,
        input_spec=tf.TensorSpec(shape=(None, 10), dtype=tf.int32),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]
    )

2. 自定义联邦训练流程，收集客户端权重 #

默认的联邦平均流程不会返回客户端完整权重，需要自定义客户端训练逻辑，让客户端返回训练后的完整权重，再在服务器端收集：

# 定义客户端训练逻辑，返回完整模型权重
@tff.tf_computation(tff_model_fn.type_signature, tff.SequenceType((tf.TensorSpec(shape=(None,10), dtype=tf.int32), tf.TensorSpec(shape=(None,), dtype=tf.int32))))
def client_train_fn(initial_weights, dataset):
    model = tff_model_fn()
    # 初始化模型权重
    tf.nest.map_structure(lambda v, t: v.assign(t), model.trainable_variables, initial_weights.trainable)
    tf.nest.map_structure(lambda v, t: v.assign(t), model.non_trainable_variables, initial_weights.non_trainable)
    
    optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
    
    # 客户端本地训练
    for x_batch, y_batch in dataset:
        with tf.GradientTape() as tape:
            logits = model(x_batch)
            loss = loss_fn(y_batch, logits)
            loss = tf.reduce_mean(loss)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    # 返回训练后的完整权重
    return tff.learning.ModelWeights(
        trainable=model.trainable_variables,
        non_trainable=model.non_trainable_variables
    )

# 初始化全局模型状态
initial_state = tff.learning.framework.get_initial_weights(tff_model_fn())

3. 每轮训练后计算权重距离 #

在联邦迭代中，每轮收集客户端权重，计算两两距离或与全局模型的距离：

# 复用问题1中的扁平化函数
def flatten_model_weights(weights):
    flat_tensors = []
    for param in weights.trainable + weights.non_trainable:
        flat_tensors.append(tf.reshape(param, shape=[-1]))
    return tf.concat(flat_tensors, axis=0)

def compute_client_weight_distances(client_weights_list):
    # 计算客户端权重两两之间的欧氏距离
    flat_weights = [flatten_model_weights(w) for w in client_weights_list]
    distance_results = []
    total_clients = len(flat_weights)
    for i in range(total_clients):
        for j in range(i+1, total_clients):
            dist = tf.norm(flat_weights[i] - flat_weights[j]).numpy()
            distance_results.append((f"客户端{i+1} vs 客户端{j+1}", dist))
    return distance_results

# 模拟联邦数据集（实际替换为真实客户端数据）
def create_synthetic_client_datasets(num_clients=3):
    client_datasets = []
    for _ in range(num_clients):
        x = tf.random.uniform((100, 10), minval=0, maxval=vocab_size, dtype=tf.int32)
        y = tf.random.uniform((100,), minval=0, maxval=vocab_size, dtype=tf.int32)
        ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(32)
        client_datasets.append(ds)
    return client_datasets

client_datasets = create_synthetic_client_datasets()

# 执行联邦迭代
global_state = initial_state
for round_num in range(10):
    # 收集所有客户端训练后的权重
    client_weights = [client_train_fn(global_state, ds) for ds in client_datasets]
    
    # 计算并打印客户端权重距离
    distances = compute_client_weight_distances(client_weights)
    print(f"=== 第{round_num+1}轮训练 ===")
    print("客户端权重两两距离：")
    for pair, dist in distances:
        print(f"{pair}: {dist:.4f}")
    
    # 更新全局模型（用联邦平均聚合客户端权重）
    client_counts = [len(list(ds)) for ds in client_datasets]
    global_state = tff.learning.federated_averaging(client_weights, client_counts=client_counts)

内容的提问来源于stack exchange，提问作者Giulio Proia