联邦学习笔记（四）：使用底层API设计联邦学习平均算法

最新推荐文章于 2024-04-12 16:33:51 发布

young_monkeysun

最新推荐文章于 2024-04-12 16:33:51 发布

阅读量3.1k

点赞数 3

分类专栏：联邦学习笔记文章标签：人工智能深度学习 python TFF

本文链接：https://blog.csdn.net/m0_46327721/article/details/109689954

版权

联邦学习笔记专栏收录该内容

6 篇文章 47 订阅

订阅专栏

写在前面

使用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值就不再下降了。说明已经达到该模型的瓶颈。换其他深度学习神经网络模型结果会更加好。

young_monkeysun

关注

3
点赞
踩
36

收藏

觉得还不错? 一键收藏
11
评论
联邦学习笔记（四）：使用底层API设计联邦学习平均算法

设计联邦学习平均算法写在前面联邦平均算法流程数据处理加载数据集数据集预处理获取处理后的数据集前向运算和损失函数获取代码输入类型创建初始模型联邦学习训练和梯度下降单batch梯度下降整个数据集梯度下降计算联邦学习评估函数本地模型评估联邦学习全局评估联邦平均算法训练实验结果loss值总结写在前面使用TFF框架实现联邦学习算法，与tf有很大差别。由于tff框架设计之初是为了能够在现实设备中部署，所以底层设计语言用的不是python。因此，使用python或者tf代码编写的算法，不能直接在TFF框架中使用。这是
复制链接

扫一扫