设计联邦学习平均算法
写在前面
使用TFF框架实现联邦学习算法,与tf有很大差别。由于tff框架设计之初是为了能够在现实设备中部署,所以底层设计语言用的不是python。因此,使用python或者tf代码编写的算法,不能直接在TFF框架中使用。这是学习TFF的关键点(把TFF当成一门新语言学习或许会更好)。
联邦平均算法流程
导入相应的包,并设置相应的环境
import tensorflow.keras as keras
from tensorflow.keras import layers , Sequential , losses , metrics , optimizers , datasets
import nest_asyncio
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus'] = False
nest_asyncio.apply()
tff.backends.reference.set_reference_context()
数据处理
实验是用的mnist
数据集,为了测试联邦平均算法在Non-IID数据集下的效果需要先将数据集切分成Non-IID类型。
假设场景 有十个客户端,每个客户端贡献自己独有的信息进行联邦学习训练
加载数据集
代码:
mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()
测试:
#数据查看
for i , d in enumerate( mnist_train[1] ):
if i <10:
print(i , d , '\n')
print([(x.dtype , x.shape) for x in mnist_train])
- 显示联邦学习数据类型和形状
- 6万张图片,每张图片像素为28*28 , 张量类型(60000 , 28 , 28)
- 6万张图片标签 (60000,)
- 以上都在单张图片中的第一维度
数据集预处理
由于实验的是图片分类任务,所以需要对图片数据进行展平,归一化处理。同时为了测试Non-IID
数据 , 需要将数据集处理成Non-IID数据。
将原始Mnist数据转换成联邦序列数据,以便联邦学习计算。mnist数据集中拥有70000张图片,其中60000张用于训练集,10000张用于测试集
数据集处理实现功能:
- 处理成Non-IID数据集
- 数据集展平,标准化
代码:
def get_data_for_digit(source, digit):
output_sequence = []
#取出相同数据标签类别的数据索引,作为一个客户端的数据,
# 比如客户端0所拥有的数据类别都是0 ,以此类推
all_samples = [i for i, d in enumerate(source[1]) if d == digit]
#将客户端数据切分成batch_size==100
for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):
batch_samples = all_samples[i:i + BATCH_SIZE]
output_sequence.append({
'x':
#将该batch_size中的数据展平,并进行(0,1)标准化
np.array([source[0][i].flatten()/255.0 for i in batch_samples],
dtype=np.float32),
'y':
#取出该batch_size中数据对应的标签
np.array([source[1][i] for i in batch_samples], dtype=np.int32)
})
return output_sequence
获取处理后的数据集
- 设置每个客户端中样本数量
- 在该示例中,所有客户端共拥有10000张图片 ,
- 每个客户端1000张,测试集共10000张图片,也就是说没有取所有的数据进行训练
代码:
NUM_EXAMPLES_PER_USER = 1000
#设置客户端训练时的batch_size,也就是说客户端一个epoch中一次训练100张图片,10次训练完
BATCH_SIZE = 100
federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]
federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]
测试:
查看客户端5中的数据,有上述可知,其中保存的应当全是标签为5的图片。
#客户端5中最后一个batch数据集的数量
print(len(federated_train_data[5][-1]['y']))
#客户端5中最后一个batch中最后一个图片数据
plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')
plt.grid(False)
plt.show()
第一个输出数据量为:
前向运算和损失函数
本实验使用的模型为全连接模型:只有一个隐藏层,隐藏层神经元为10个,输出层神经元1个。
代码:
@tf.function
def forward_pass(model , batch):
# CC_loss = losses.CategoricalCrossentropy(from_logits=True)
pre_y = tf.nn.softmax(
tf.matmul(batch['x'], model['weights']) + model['bias']
)
return -tf.reduce_mean(
# CC_loss(tf.one_hot(batch['y'], depth = 10 , on_value = None) , pre_y)
tf.reduce_sum(
tf.one_hot(batch['y'], 10 , on_value=None , off_value=None ) * tf.math.log(pre_y), axis=[1])
)
@tff.tf_computation(MODEL_TYPE , BATCH_TYPE)
def batch_loss(model , batch):
return forward_pass(model , batch)
注意:
- 由于TFF底层的代码不是python,所以不能使用python类型直接进行计算。使用python类型需要对其进行tff装饰,装饰成tff可以使用的类型
- 使用独立的python函数编写复杂的TF逻辑代码,不推荐使用tff.tf_computation装饰器
- 在装饰器tff.tf_computation中使用tf.function装饰器 ,则可以在其内部使用python语言编写代码逻辑。tff.tf_computation 函数可以调用tf.function装饰器 装饰的代码逻辑,但是反之不能。
也就是说在TFF框架中不能直接使用python代码,需要先将python装饰成tf代码,再将tf代码装饰成TFF代码。
获取代码输入类型
在batch_loss
函数中需要提供输入参数的数据类型(MODEL_TYPE , BATCH-TYPE
)。
代码:
#定义联邦训练时的输入规格
BATCH_SPEC = collections.OrderedDict(
#因为一个batch中数据的数量不确定,所以定义为未知
x = tf.TensorSpec(shape=[None , 784] , dtype=tf.float32) ,
#由于标签类型在内存中存储的是int类型,所以需要定义成int类型,否则会报错
y = tf.TensorSpec(shape=[None] , dtype=tf.int32)
)
#batch的输入类型
BATCH_TYPE = tff.to_type(BATCH_SPEC)
#获取神经网络模型的参数规格
MODEL_SPEC = collections.OrderedDict(
#10表示输入层有10个神经元
weights = tf.TensorSpec(shape=[784 , 10],dtype=tf.float32),
bias = tf.TensorSpec(shape=[10] , dtype=tf.float32)
)
#网络模型的输入类型
MODEL_TYPE = tff.to_type(MODEL_SPEC)
创建初始模型
代码:
initial_model = collections.OrderedDict(
weights = np.zeros([784,10] , dtype=np.float32),
bias = np.zeros([10] , dtype=np.float32)
)
测试:
测试上述计算loss值函数
#标签为5的最后一个batch数据集进行测试
client_5_last_batch = federated_train_data[5][-1]
print( batch_loss(initial_model , client_5_last_batch) )
测试结果:
联邦学习训练和梯度下降
单batch梯度下降
- 由于序列化会丢失一些调试信息,同时难以调试,所以尽量不要使用tff.tf_computation装饰器来debug tf代码
- tff.tf_computation装饰的函数可以内联在tff.tf_computation装饰的函数中,但是不建议这样做。
- 在tff.tf_computation装饰的函数中最好将内部函数设置为常规的python函数或者tff函数
- 如果需要调用的函数是tff类型的,则要使用tff装饰
由于训练是逐个batch进行训练的,所以实现在单个batch进行梯度下降的函数。
@tff.tf_computation(MODEL_TYPE , BATCH_TYPE , tf.float32)
def batch_train(initial_model , batch , lr):
model_vars = collections.OrderedDict([
(name , tf.Variable(name=name , initial_value=value))
for name , value in initial_model.items()
])
optimizer = optimizers.SGD(learning_rate= lr)
@tf.function
def _train_on_batch(model_vars , batch):
#通过由上述函数求得的loss来对模型进行梯度下降
with tf.GradientTape() as tape:
loss = forward_pass(model_vars , batch)
grads = tape.gradient(loss , model_vars)
optimizer.apply_gradients(zip(tf.nest.flatten(grads) , tf.nest.flatten(model_vars) ))
return model_vars
return _train_on_batch(model_vars , batch)
再次强调:
在有关TFF计算中只能使用TFF类型。逻辑如下图
测试代码:
print(str(batch_train.type_signature))
#梯度下降测试
model = initial_model
sum_loss = []
for _ in range(5):
model = batch_train(model , client_5_last_batch , 0.01)
sum_loss.append(batch_loss(model , client_5_last_batch))
print(sum_loss)
测试结果:
整个数据集梯度下降计算
在客户端中通过调用上述batch_train
函数实现整个数据集的训练。
由于是对整个数据集进行操作,所以首先获得整个数据集的类型。整个数据集是被列表化的,可通过ttf.SequenceType(BATCH_TYPE)
获取整个数据集的类型。
代码:
#定义本地客户端输入所有数据的数据类型
LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)
#定义该函数的数据输入类型
@tff.federated_computation(MODEL_TYPE , tf.float32 , LOCAL_DATA_TYPE)
def local_train(initial_model , lr , all_batchs):
#每个batch数据都使用MAP函数进行计算
@tff.federated_computation(MODEL_TYPE , BATCH_TYPE)
def batch_fn(model , batch):
#由于batch_train函数接收的是三个参数,而tff.sequence_reduce处理的是两个参数
#因此将batch_train嵌入到该函数中,由外函数提供学习率lr
return batch_train(model , batch , lr)
#模型训练,将所有客户端数据,逐个batch的调用上述的batch_train函数进行训练
#使用的是batch_train中的SGD梯度下降方法,
# 当对所有batch数据训练完成,也就对整个客户端数据训练完成
#tff.sequence_reduce函数应用于联邦计算函数中,tff.sequence_reduce内不能包含tf代码
return tff.sequence_reduce(all_batchs , initial_model , batch_fn)
测试:
查看该函数的状态
print(local_train.type_signature)
结果:
联邦学习评估函数
本地模型评估
在评估本地模型时,有两种策略:
- 计算本地所有batch数据的总loss值,与初始模型状态对比看是否下降
- 计算本地所有batch数据的平均loss值,与初始模型状态相比看是否下降
代码:
@tff.federated_computation(MODEL_TYPE , LOCAL_DATA_TYPE)
def local_eval(model , all_batchs):
#计算平均loss值时使用tff.sequence_average函数
return tff.sequence_sum(
#序列map函数,将all_batchs中的每个batch都进行loss值计算
#tff.sequence_map与tff.sequence_reduce的差别在于map是并行计算,reduce是串行计算
tff.sequence_map(
tff.federated_computation(lambda b: batch_loss(model , b) ,BATCH_TYPE) ,
all_batchs
)
)
测试:
print(local_eval.type_signature)
#评估初始模型在客户端5上的表现
#在客户端5中训练所有数据
client_5_batchs = federated_train_data[5]
client_5_train_model = local_train(initial_model , 0.01 , client_5_batchs)
print('init model on client_5_datasets loss = ' , local_eval(initial_model , client_5_batchs) , '\n')
print('client5 model on client_5_datasets loss = ' , local_eval(client_5_train_model , client_5_batchs) , '\n')
client_0_batchs = federated_train_data[0]
print('client5 model on client_0_datasets loss = ' , local_eval(client_5_train_model , client_0_batchs) , '\n')
测试结果:
由测试结果可知,初始模型在客户端5中数据集上的效果不佳。当初始模型在客户端5中的数据集进行训练后,其loss值显著减小。但是初始模型在客户端5中的数据集进行训练后获得的模型在客户端0中数据集的表现不佳。因为这是Non-IID类型的数据。
联邦学习全局评估
进行联邦学习全局评估:
- 获取服务器端的模型类型
- 获取客户端的数据类型
代码:
SERVER_MODEL_TYPE = tff.type_at_server(MODEL_TYPE)
CLIENT_DATA_TYPE = tff.type_at_clients(LOCAL_DATA_TYPE)
全局评估函数:
- 将聚合的模型分发给所有的被选中的客户端,
- 所有客户端接收到全局模型后,使用本地数据进行计算,求取本地的平均损失率,用于评估联邦模型的效果
- 模型服务端发送到客户端
- 平均loss值客户端发送到服务端
代码:
@tff.federated_computation(SERVER_MODEL_TYPE , CLIENT_DATA_TYPE)
def federated_eval(model , data):
return tff.federated_mean(
#[]是隐式转换到tff类型
tff.federated_map(local_eval , [tff.federated_broadcast(model) , data] )
)
#在tff中存在隐式转换,例如{<X,Y>}@Z 等价于 <{X}@Z , {Y}@Z>
测试:
计算上述在客户端5中产生的模型其他所有客户端中的loss值 ,求所有客户端的平均loss值作为衡量联邦学习效果。
print(federated_eval.type_signature)
print('initial model federated loss = ',federated_eval(initial_model , federated_train_data),'\n')
print('client_5_train_model federated loss = ' , federated_eval(client_5_train_model , federated_train_data),'\n')
测试结果:
由测试结果可知,虽然客户端5的训练模型在本地训练效果很好,但是在全局其loss值反而增大了。
联邦平均算法
联邦平均算法:客户端在本地训练完成之后,上传本地模型参数到服务端,服务端再将所有收集的参数值求平均,获得全局模型参数。
代码:
#将tf.float32类型封装成tff中的float32类型
SERVER_FLOAT_TYPE = tff.type_at_server(tf.float32)
@tff.federated_computation(SERVER_MODEL_TYPE , SERVER_FLOAT_TYPE , CLIENT_DATA_TYPE)
def federated_train(model , lr ,data):
#返回的是训练后的模型
return tff.federated_mean(
tff.federated_map(local_train , [
tff.federated_broadcast(model) ,
tff.federated_broadcast(lr),
data
])
)
测试:
print(federated_train.type_signature)
测试结果:
训练
代码:
#联邦训练测试
model = initial_model
lr = 0.1
for round in range(5):
model = federated_train(model, lr , federated_train_data)
lr = lr*0.9
loss = federated_eval(model , federated_train_data)
print('round {} , loss = {}'.format(round , loss))
#查看模型在测试集上的效果
print('initial_model test loss =',
federated_eval(initial_model, federated_test_data))
print('trained_model test loss =', federated_eval(model, federated_test_data))
实验结果
loss值
初始模型在测试集上的loss值和训练后的模型在测试集上的loss值。
总结
上述实验结果显示,在训练第80轮左右,loss值就不再下降了。说明已经达到该模型的瓶颈。换其他深度学习神经网络模型结果会更加好。