# 导入库
import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
# 定义卷积神经网络模型
def create_model():
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.MaxPooling2D((2, 2)),
tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
tf.keras.layers.MaxPooling2D((2, 2)),
tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
return model
# 定义联邦学习FedAvg算法
def federated_averaging(num_clients, epochs, batch_size, lr):
# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 对数据进行预处理
x_train = x_train.reshape((-1, 28, 28, 1)).astype(np.float32) / 255.0
x_test = x_test.reshape((-1, 28, 28, 1)).astype(np.float32) / 255.0
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)
# 定义全局模型
global_model = create_model()
# 复制全局模型作为本地模型
local_models = [tf.keras.models.clone_model(global_model) for _ in range(num_clients)]
# 定义优化器
optimizer = tf.keras.optimizers.Adam(lr=lr)
# 定义损失函数
loss_fn = tf.keras.losses.CategoricalCrossentropy()
# 进行联邦学习
for epoch in range(epochs):
# 在每个客户端上训练本地模型
for i in range(num_clients):
# 获取本地训练数据
local_x_train, local_y_train = x_train[i*batch_size:(i+1)*batch_size], y_train[i*batch_size:(i+1)*batch_size]
# 在本地模型上进行训练
local_models[i].compile(optimizer=optimizer, loss=loss_fn, metrics=['accuracy'])
local_models[i].fit(local_x_train, local_y_train, epochs=1, verbose=0)
# 对本地模型进行聚合
for layer in global_model.layers:
if isinstance(layer, tf.keras.layers.Conv2D) or isinstance(layer, tf.keras.layers.Dense):
# 获取本地模型的参数
local_params = [local_model.get_weights()[i] for local_model in local_models for i in range(len(layer.get_weights()))]
# 将本地模型的参数进行平均
global_params = np.mean(local_params, axis=0)
# 将全局模型的参数更新为平均值
layer.set_weights(global_params)
# 在测试集上测试全局模型
global_model.compile(optimizer=optimizer, loss=loss_fn, metrics=['accuracy'])
loss, accuracy = global_model.evaluate(x_test, y_test, verbose=0)
print('Epoch', epoch, 'Test loss:', loss, 'Test accuracy:', accuracy)
# 运行联邦学习FedAvg算法
federated_averaging(num_clients=10, epochs=10, batch_size=32, lr=0.001)
订阅专栏 解锁全文
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
2264
