联邦学习系列文章:
本篇文章与上篇《【联邦学习】用Tensorflow实现联邦模型AlexNet on CIFAR-10》类似,也是用tensorflow在单机下模拟联邦的过程,因此有些描述会比较简略,建议先看完上一篇。与之不同的点总结如下:
对于一个神经网络模型来说,训练的目标是使得模型输出与监督信息(或自监督信息)尽量相近。这个想法通常体现在loss函数的设计上。例如,对于二分类问题,神经网络的输出值通常为0~1
之间的浮点数(例如0.31
),而监督信息通常为{0, 1}
这样的离散值(例如0
)。如何判断0.3
与0
之间是否“相近”,通常使用“均方误差”、“均方根误差”等loss函数来计算,最终得到一个标量值,也就是我们常说的loss值。在数学上,这个loss函数可以定义为:
L ( w , S ) , L(w, S), L(w,S),
其中, w w w是模型需要训练的参数, S = ( X , Y ) S=(X, Y) S=(X,Y)是训练数据集。有时loss函数也会等价地定义为 L ( f w ( X ) , Y ) L(f_w(X), Y) L(fw(X),Y)。那么,对于每一次梯度下降来说,所谓的梯度即是loss函数对模型参数的逐一求导:
g i = ∂ L ( w , S ) ∂ w i . g_i = \frac{\partial L(w, S)}{\partial w_i}. gi=∂wi∂L(w,S).
可以看出,对于每一个参数都会有一个标量梯度值,因此整体梯度的大小是与模型参数的大小一致的。例如,我们要训练一个双层的DNN,一共有300个权重参数,那么经过一轮反向传播求导后,得到的梯度也会有300个数值。
一行印在所有炼丹人DNA里文字:“负梯度方向是(loss函数)下降最快的方向”。梯度下降方法(可能)是神经网络参数更新的唯一方法,基于上面求导得到的梯度,参数更新过程可以表示为:
w i ( t + 1 ) = w i ( t ) − ρ g i , w_i^{(t+1)} = w_i^{(t)} - \rho~g_i, wi(t+1)=wi(t)−ρ gi,
其中, t t t是迭代次数, ρ \rho ρ是学习率。与最优化领域的梯度下降不同,神经网络训练中的学习率通常是预定的超参数,而不需要计算。
对于多轮梯度下降:
w i ( t + 1 ) = w i ( t ) − ρ g i ( t ) , w i ( t ) = w i ( t − 1 ) − ρ g i ( t − 1 ) , . . . w i ( 1 ) = w i ( 0 ) − ρ g i ( 0 ) . w_i^{(t+1)} = w_i^{(t)} - \rho~g^{(t)}_i,\\ w_i^{(t)} = w_i^{(t-1)} - \rho~g^{(t-1)}_i,\\ ...\\ w_i^{(1)} = w_i^{(0)} - \rho~g^{(0)}_i. wi(t+1)=wi(t)−ρ gi(t),wi(t)=wi(t−1)−ρ gi(t−1),...wi(1)=wi(0)−ρ gi(0).
其实可以简单得到:
w i ( t + 1 ) = w i ( 0 ) − ρ ∑ k t g i ( k ) . w_i^{(t+1)} = w_i^{(0)} - \rho~\sum_k^t g^{(k)}_i. wi(t+1)=wi(0)−ρ k∑tgi(k).
对于训练过程的一个epoch,通常包含多个mini-batch梯度下降。根据上式,这些batch的梯度求和起来,可以表示为一个epoch的总梯度。这个概念会在联邦过程中重点使用。
回顾一下联邦学习框架中的两个角色:
我们把联邦训练的流程整理成一张流程图,包括一个Server和两个Clients。
图中1、2、4步都很容易实现,而第3步传输的“模型更新”具体是什么,是值得探究的。在前一篇文章中,笔者提到使用tf.keras.optimizers.Optimizer().get_gradients()
来获取梯度 g i g_i gi,其实是有一定的误导性的。因为普通的梯度下降(GD)是把梯度直接应用到参数上,即:
w i ( t + 1 ) = w i ( t ) − ρ g i . w_i^{(t+1)} = w_i^{(t)} - \rho~g_i. wi(t+1)=wi(t)−ρ gi.
可以使用get_gradients()
的结果来作为模型更新梯度。而其他启发式优化器(如Adam)会将梯度 g i g_i gi做一次修改,加上一些启发式的信息得到新的梯度 g i ^ = A d a m ( g i ) \hat{g_i} = Adam(g_i) gi^=Adam(gi),再应用到参数上,即:
w i ( t + 1 ) = w i ( t ) − ρ g i ^ . w_i^{(t+1)} = w_i^{(t)} - \rho~\hat{g_i}. wi(t+1)=wi(t)−ρ gi^.
这时如果只记录get_gradients()
计算的梯度发送给Server,其实是错误的。那么有什么办法获取到Adam()
输出的值吗?目前笔者没有找到相关接口。但是,只要把上式做一个简单的变换,就可以得到:
g i ^ = ( w i ( t ) − w i ( t + 1 ) ) / ρ . \hat{g_i} = (w_i^{(t)} - w_i^{(t+1)}) / \rho. gi^=(wi(t)−wi(t+1))/ρ.
即我们只需要知道更新前后的参数值,就可以得到优化器输出的梯度值。对于 t t t个mini-batch SGD组成的一个epoch,也可以通过类似的方法得到该epoch的整体梯度:
∑ k t g i ( k ) = ( w i ( 0 ) − w i ( t + 1 ) ) / ρ . \sum_k^t g^{(k)}_i = (w_i^{(0)} - w_i^{(t+1)}) / \rho~. k∑tgi(k)=(wi(0)−wi(t+1))/ρ .
等式的左边即是Client端执行完一轮模型更新(可能包含多个epoch)后,需要发送给Server端的模型更新梯度。体现在代码实现中,每个Client就只需要记录一下刚收到的全局模型参数,以及更新后的模型参数,代入到上式中即可获得该Client的本地模型更新梯度。
本次代码实现使用了一个最简单的线性回归模型:
y = W x + b y = Wx+b y=Wx+b
在一个简单的二分类数据集上进行测试(暂时忘记数据集出处了,之后补上),所以代码部分不再赘述太多。有TF基础的同学可以点击这里直接查看jupyter版(src_grad
目录),数据集也在该仓库的src_grad/data
文件夹中。
读取数据集,按照Client的数量设置划分成多份训练集、以及一份全局的测试集。
from __future__ import print_function, division
import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import random
def split_data(path, clients_num):
# 读取数据
data = pd.read_csv(path)
# 拆分数据
X_train, X_test, y_train, y_test = train_test_split(
data[["Temperature", "Humidity", "Light", "CO2", "HumidityRatio"]].values,
data["Occupancy"].values.reshape(-1, 1),
random_state=42)
# one-hot 编码
y_train = np.concatenate([1 - y_train, y_train], 1)
y_test = np.concatenate([1 - y_test, y_test], 1)
# 训练集划分给多个client
X_train = np.array_split(X_train, clients_num)
y_train = np.array_split(y_train, clients_num)
return X_train, X_test, y_train, y_test
CLIENT_NUM = 6
X_train, X_test, y_train, y_test = split_data("./data/datatraining.txt", CLIENT_NUM)
使用文件系统来模拟网络传输,也可以接入区块链等分布式存储方式。主要包含四个功能:
import os
import pickle
import gzip
BASE_DIR = "./storage"
if not os.path.isdir(BASE_DIR):
os.mkdir(BASE_DIR)
def pack(model):
pkl = pickle.dumps(model)
pkl = gzip.compress(pkl)
return pkl
def unpack(data):
pkl = gzip.decompress(data)
model = pickle.loads(pkl)
return model
def client_query_model():
"""return the newest model and epoch num"""
newest_epoch = -1
res_f = None
for f in os.listdir(BASE_DIR):
if not f.startswith('global_model'):
continue
file_name = os.path.splitext(f)[0]
epoch = int(file_name.split('_')[-1])
if epoch > newest_epoch:
newest_epoch = epoch
res_f = f
# file found
with open("{}/{}".format(BASE_DIR, res_f), 'rb') as rf:
res = rf.read()
return unpack(res), newest_epoch
def client_upload_one_update(update, epoch, c_id):
"""upload one model update"""
file_name = "{}/local_update_{}_{}.ieen".format(BASE_DIR, c_id, epoch)
data = pack(update)
with open(file_name, 'wb') as wf:
wf.write(data)
return
def server_query_updates(cur_epoch):
"""query all model updates"""
res = []
for f in os.listdir(BASE_DIR):
if not f.startswith('local_update'):
continue
file_name = os.path.splitext(f)[0]
epoch = int(file_name.split('_')[-1])
if epoch == cur_epoch:
with open("{}/{}".format(BASE_DIR, f), 'rb') as rf:
data = unpack(rf.read())
res.append(data)
return res
def server_upload_model(model, epoch):
"""upload one model with epoch num"""
file_name = "{}/global_model_{}.ieen".format(BASE_DIR, epoch)
data = pack(model)
with open(file_name, 'wb') as wf:
wf.write(data)
return
Client获取到全局模型后,使用全局模型的参数来初始化本地模型的参数,之后启动mini-batch SGD,最后计算参数更新梯度,发送给Server。
# client 要训练的epoch
client_epoch = [0] * CLIENT_NUM
client_learning_rate = 0.001
def train_model(client_id):
model, epoch = client_query_model()
if epoch < client_epoch[client_id]:
return
tf.compat.v1.reset_default_graph()
n_samples = X_train[client_id].shape[0]
x = tf.placeholder(tf.float32, [None, n_features])
y = tf.placeholder(tf.float32, [None, n_class])
ser_W, ser_b = model
W = tf.Variable(ser_W)
b = tf.Variable(ser_b)
pred = tf.matmul(x, W) + b
# 定义损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=pred,
labels=y))
# 梯度下降
# optimizer = tf.train.AdamOptimizer(learning_rate)
optimizer = tf.train.GradientDescentOptimizer(client_learning_rate)
gradient = optimizer.compute_gradients(cost)
train_op = optimizer.apply_gradients(gradient)
# 初始化所有变量
init = tf.global_variables_initializer()
# 训练模型
with tf.Session() as sess:
sess.run(init)
avg_cost = 0
total_batch = int(n_samples / batch_size)
for i in range(total_batch):
_, c = sess.run(
[train_op, cost],
feed_dict={
x: X_train[client_id][i * batch_size:(i + 1) * batch_size],
y: y_train[client_id][i * batch_size:(i + 1) * batch_size, :]
})
avg_cost += c / total_batch
# 获取更新量
val_W, val_b = sess.run([W, b])
delta_W = (ser_W-val_W)/client_learning_rate
delta_b = (ser_b-val_b)/client_learning_rate
delta_model = [delta_W, delta_b]
meta = [n_samples, avg_cost]
client_upload_one_update([delta_model, meta], epoch, client_id)
client_epoch[client_id] = epoch
return
Server端初始化一个全局的模型参数,并开始(串行地)调度各个Client进行训练,然后聚合它们发回的模型更新梯度,以更新全局参数。每轮都跑一下测试集,看看训练效果。
# 跑测试集
def testing(ser_W, ser_b):
tf.compat.v1.reset_default_graph()
x = tf.placeholder(tf.float32, [None, n_features])
y = tf.placeholder(tf.float32, [None, n_class])
W = tf.Variable(ser_W)
b = tf.Variable(ser_b)
pred = tf.matmul(x, W) + b
correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# 初始化所有变量
init = tf.global_variables_initializer()
# 跑模型
with tf.Session() as sess:
sess.run(init)
acc = accuracy.eval({x: X_test, y: y_test})
return acc
# 设置模型
batch_size = 100
n_features = 5
n_class = 2
EPOCH_NUM = 50 * CLIENT_NUM
server_lr = 0.001
# 模型参数
server_W = np.zeros([n_features, n_class], dtype=np.float32)
server_b = np.zeros([n_class], dtype=np.float32)
server_model = [server_W, server_b]
for epoch in range(EPOCH_NUM):
server_upload_model(server_model, epoch)
for c_id in range(CLIENT_NUM):
train_model(c_id)
total_grad_W = None
total_grad_b = None
total_size = 0
total_cost = 0
updates = server_query_updates(epoch)
for update in updates:
grads, meta = update
grad_W, grad_b = grads
data_size, cost = meta
total_grad_W = (grad_W * data_size) if (total_grad_W is None) else (total_grad_W + grad_W * data_size)
total_grad_b = (grad_b * data_size) if (total_grad_b is None) else (total_grad_b + grad_b * data_size)
total_size += data_size
total_cost += cost
total_grad_W /= total_size
total_grad_b /= total_size
total_cost /= CLIENT_NUM
# update global model
server_W = server_W - server_lr * total_grad_W
server_b = server_b - server_lr * total_grad_b
server_model = [server_W, server_b]
test_acc = testing(server_W, server_b)
print("Epoch: {:03}, cost: {:.2f}, test_acc: {:.4f}".format(epoch, total_cost, test_acc))
当网络比较复杂时,可以使用trainable_variables()
函数获取所有的可训练的参数列表。当网络结构固定后,这个列表内的变量顺序不会改变。
上文提到,使用了Adam等优化器时,Client发送的其实是模型参数更新增量,而优化器中的“历史梯度信息”就被丢弃了。如何解决这个问题,已经有相关论文进行了讨论,大概想法是把优化器中的参数也一起联邦传输。具体如何实现请读者参考链接中的文章。