联邦学习实战：用python从零实现横向联邦图像分类
本文主要参考了杨强教授的《联邦学习实战》
https://github.com/FederatedAI/Practicing-Federated-Learning/tree/main/chapter03_Python_image_classification
一、环境配置
python3.7
GPU（可选）：首先安装CUDA、cuDNN
安装pytorch，pip install torch
二、python实现横向联邦图像分类
2.1配置信息
- 训练的客户端数量 no_models
- 全局迭代次数 global_epochs
- 本地模型的迭代次数 local_epochs
- 本地训练相关的算法配置，lr（学习率）等
- 模型信息：本案例用的是ResNet-18图像分类模型
- 数据信息：本案例使用的是cifar10数据集
其他配置信息，例如是否使用差分隐私，模型聚合策略，可以根据需求自行添加。这里将上面的信息以json格式记录在配置文件中，以便修改。文件命名为conf.json
{
"model_name" : "resnet18",
"no_models" : 10,
"type" : "cifar",
"global_epochs" : 20,
"local_epochs" : 3,
"k" : 6,
"batch_size" : 32,
"lr" : 0.001,
"momentum" : 0.0001,
"lambda" : 0.1
}
上面的k指的是每一轮迭代时，服务端会从所有客户端中挑选k个客户端参与训练。
联邦学习在模型训练之前，会将配置信息发送到客户端和服务端保存，如果配置信息更改，也会同时对所有参与方同步。
2.2训练数据集
这里使用的是torchvision的datasets模块内置的cifar10数据集， 项目文件下创建getDatasets.py文件
from torchvision.transforms import transforms
from torchvision import datasets
# 获取数据集
def get_dataset(dir, name):
if name == 'mnist':
# root: 数据路径
# train参数表示是否是训练集或者测试集
# download=true表示从互联网上下载数据集并把数据集放在root路径中
# transform：图像类型的转换
train_dataset = datasets.MNIST(dir, train=True, download=True, transform=transforms.ToTensor())
eval_dataset = datasets.MNIST(dir, train=False, transform=transforms.ToTensor())
elif name == 'cifar':
# 设置两个转换格式
# transforms.Compose 是将多个transform组合起来使用（由transform构成的列表）
transform_train = transforms.Compose([
# transforms.RandomCrop： 切割中心点的位置随机选取
transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
# transforms.Normalize： 给定均值：(R,G,B) 方差：（R，G，B），将会把Tensor正则化
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
train_dataset = datasets.CIFAR10(dir, train=True, download=True, transform=transform_train)
eval_dataset = datasets.CIFAR10(dir, train=False, transform=transform_test)
return train_dataset, eval_dataset
2.3服务端
服务端类Server主要包括三种函数
- 定义构造函数
- 模型聚合函数
- 模型评估函数
2.3.1定义构造函数
class Server(object):
# 定义构造函数
def __init__(self, conf, eval_dataset):
# 导入配置文件
self.conf = conf
# 根据配置获取模型文件
self.global_model = models.get_model(self.conf["model_name"])
self.eval_loader = torch.utils.data.DataLoader(
eval_dataset,
# 设置单个批次大小
batch_size=self.conf["batch_size"],
# 打乱数据集
shuffle=True
)
2.3.2模型聚合函数
这里使用的是经典的FedAvg算法
def model_aggregate(self, weight_accumulator):
# weight_accumulatot存储了每一个客户端的上传参数变化值
# 遍历服务器的全局模型
for name, data in self.global_model.state_dict().items():
# 更新每一层
update_per_layer = weight_accumulator[name] * self.conf["lambda"]
# 累加和
if data.type() != update_per_layer.type():
# 因为update_per_layer的type是floatTensor，所以将起转换为模型的LongTensor（有一定的精度损失）
data.add_(update_per_layer.to(torch.int64))
else:
data.add_(update_per_layer)
2.3.3 模型评估函数
对当前的全局模型，利用评估数据评估当前的全局模型性能
def model_eval(self):
self.global_model.eval()
total_loss = 0.0
correct = 0
dataset_size = 0
# 遍历评估数据集合
for batch_id, batch in enumerate(self.eval_loader):
data, target = batch
# 获取所有的样本总量大小
dataset_size += data.size()[0]
# 将数据存储到gpu
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 加载到模型中训练
output = self.global_model(data)
# 聚合所有的损失 cross_entropy交叉熵函数计算损失
total_loss += torch.nn.functional.cross_entropy(output,target,
reduction='sum'
).item()
# 获取最大的对数概率的索引值，在预测结果中选择可能性最大的作为最终的分类结果
pred = output.data.max(1)[1]
# 统计预测结果与真实标签相同的总个数
correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
# 计算准确率
acc = 100.0 * (float(correct) / float(dataset_size))
# 计算损失值
total_1 = total_loss / dataset_size
return acc, total_1
服务端汇总到项目文件夹下的server.py文件，如下
import torch
import models
class Server():
def __init__(self, conf, eval_dataset):
# 导入配置文件
self.conf = conf
# 根据配置获取模型文件
self.global_model = models.get_model(self.conf["model_name"])
self.eval_loader = torch.utils.data.DataLoader(
eval_dataset,
# 设置单个批次大小
batch_size=self.conf["batch_size"],
# 打乱数据集
shuffle=True
)
def model_aggregate(self, weight_accumulator):
# weight_accumulatot存储了每一个客户端的上传参数变化值
# 遍历服务器的全局模型
for name, data in self.global_model.state_dict().items():
# 更新每一层
update_per_layer = weight_accumulator[name] * self.conf["lambda"]
# 累加和
if data.type() != update_per_layer.type():
# 因为update_per_layer的type是floatTensor，所以将起转换为模型的LongTensor（有一定的精度损失）
data.add_(update_per_layer.to(torch.int64))
else:
data.add_(update_per_layer)
def model_eval(self):
self.global_model.eval()
total_loss = 0.0
correct = 0
dataset_size = 0
# 遍历评估数据集合
for batch_id, batch in enumerate(self.eval_loader):
data, target = batch
# 获取所有的样本总量大小
dataset_size += data.size()[0]
# 将数据存储到gpu
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 加载到模型中训练
output = self.global_model(data)
# 聚合所有的损失 cross_entropy交叉熵函数计算损失
total_loss += torch.nn.functional.cross_entropy(
output, target, reduction='sum').item()
# 获取最大的对数概率的索引值，在预测结果中选择可能性最大的作为最终的分类结果
pred = output.data.max(1)[1]
# 统计预测结果与真实标签相同的总个数
correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
# 计算准确率
acc = 100.0 * (float(correct) / float(dataset_size))
# 计算损失值
total_1 = total_loss / dataset_size
return acc, total_1
2.4 客户端
客户端类包括两种函数
- 构造函数
- 本地训练函数
客户端主要工作包括：
首先，将配置信息复制到客户端中；
然后，按照配置中的模型信息获取模型，通常由服务端将模型参数传递给客户端，客户端将该全局模型覆盖掉本地模型；
最后，配置本地训练数据
这里通过torchvision中的datasets模块获取cifar10数据集后，按客户端ID进行切分，不同客户端拥有不同的子数据集，相互之间没有交集
2.4.1 定义构造函数
class Client(object):
def __init__(self, conf, model, train_dataset, id = 1):
# 配置文件
self.conf = conf
# 客户端本地模型
self.local_model = model
# 客户端ID
self.client_id = id
# 客户端本地数据集
self.train_dataset = train_dataset
# 按ID对训练集合的拆分
all_range = list(range(len(self.train_dataset)))
data_len = int(len(self.train_dataset) / self.conf['no_models'])
indices = all_range[id * data_len: (id + 1) * data_len]
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=conf["batch_size"],
# sampler定义从数据集中提取样本的策略
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices)
)
2.4.2 定义模型本地训练函数
这里使用交叉熵作为本地模型的损失函数，利用梯度下降求解并更新参数值
def local_train(self, model):
for name, param in model.state_dict().items():
# 客户端首先用服务器端下发的全局模型覆盖本地模型
self.local_model.state_dict()[name].copy_(param.clone())
# 定义最优化函数器，用于本地模型训练
optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'], momentum=self.conf['momentum'])
# 本地训练模型
self.local_model.train()
for e in range(self.conf["local_epochs"]):
for batch_id, batch in enumerate(self.train_loader):
data, target = batch
# 数据加载到gpu
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 梯度清0
optimizer.zero_grad()
# 训练预测
output = self.local_model(data)
# 使用cross_entropy交叉熵计算损失函数
loss = torch.nn.functional.cross_entropy(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
print("Epoch %d done" % e)
diff = dict()
for name, data in self.local_model.state_dict().items():
# 计算训练后与训练前的差值
diff[name] = (data - model.state_dict()[name])
print("Client %d local train done" % self.client_id)
return diff
客户端汇总到项目文件夹下的client.py文件，如下：
import torch
class Client():
def __init__(self, conf, model, train_dataset, id=1):
# 配置文件
self.conf = conf
# 客户端本地模型
self.local_model = model
# 客户端ID
self.client_id = id
# 客户端本地数据集
self.train_dataset = train_dataset
# 按ID对训练集合的拆分
all_range = list(range(len(self.train_dataset)))
data_len = int(len(self.train_dataset) / self.conf['no_models'])
indices = all_range[id * data_len: (id + 1) * data_len]
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=conf["batch_size"],
# sampler定义从数据集中提取样本的策略
sampler=torch.utils.data.sampler.SubsetRandomSampler(indices)
)
def local_train(self, model):
for name, param in model.state_dict().items():
# 客户端首先用服务器端下发的全局模型覆盖本地模型
self.local_model.state_dict()[name].copy_(param.clone())
# 定义最优化函数器，用于本地模型训练
optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'], momentum=self.conf['momentum'])
# 本地训练模型
self.local_model.train()
for e in range(self.conf["local_epochs"]):
for batch_id, batch in enumerate(self.train_loader):
data, target = batch
# 数据加载到gpu
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 梯度清0
optimizer.zero_grad()
# 训练预测
output = self.local_model(data)
# 使用cross_entropy交叉熵计算损失函数
loss = torch.nn.functional.cross_entropy(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
print("Epoch %d done" % e)
diff = dict()
for name, data in self.local_model.state_dict().items():
# 计算训练后与训练前的差值
diff[name] = (data - model.state_dict()[name])
print("Client %d local train done" % self.client_id)
return diff
2.5模型文件
项目文件夹下创建models.py文件，用于各种机器学习模型使用
import torch
from torchvision import models
# 各种机器学习模型
def get_model(name="vgg16", pretrained=True):
if name == "resnet18":
model = models.resnet18(pretrained=pretrained)
elif name == "resnet50":
model = models.resnet50(pretrained=pretrained)
elif name == "densenet121":
model = models.densenet121(pretrained=pretrained)
elif name == "alexnet":
model = models.alexnet(pretrained=pretrained)
elif name == "vgg16":
model = models.vgg16(pretrained=pretrained)
elif name == "vgg19":
model = models.vgg19(pretrained=pretrained)
elif name == "inception_v3":
model = models.inception_v3(pretrained=pretrained)
elif name == "googlenet":
model = models.googlenet(pretrained=pretrained)
if torch.cuda.is_available():
return model.cuda()
else:
return model
2.6 整合
首先，读取配置文件信息
# 读取配置文件
with open(args.conf, 'r') as f:
conf = json.load(f)
接下来，分别定义一个服务端对象和多个客户端对象，模拟横向联邦学习场景
train_datasets, eval_datasets = datasets.get_dataset("./data/", conf["type"])
server = Server(conf, eval_datasets)
# 客户端列表
clients = []
# 添加10个客户端到列表
for c in range(conf["no_models"]):
clients.append(Client(conf, server.global_model, train_datasets, c))
每一轮迭代，服务器会从当前客户端中随机挑选一部分参与本轮训练，被选中的客户端用本地训练接口local_train进行本地训练，最后服务端调用模型聚合函数更新全局模型
for e in range(conf["global_epochs"]):
print("Global Epoch %d" % e)
# 随机采样k个客户端参与本轮联邦训练
candidates = random.sample(clients, conf["k"])
print("select clients is: ")
for c in candidates:
print(c.client_id)
weight_accumulator = {}
# 初始化模型参数weight_accumulator
for name, params in server.global_model.state_dict().items():
# 生成一个和参数矩阵大小相同的0矩阵
weight_accumulator[name] = torch.zeros_like(params)
# 遍历被挑选的客户端，本地训练模型
for c in candidates:
diff = c.local_train(server.global_model)
for name, params in server.global_model.state_dict().items():
weight_accumulator[name].add_(diff[name])
# 模型聚合
server.model_aggregate(weight_accumulator)
# 模型评估
acc, loss = server.model_eval()
print("Epoch %d, acc: %f, loss: %f\n" % (e, acc, loss))
项目文件夹下创建main.py文件，用来整合代码，如下：
import argparse
import json
import random
import getDatasets
from client import *
from server import *
if __name__ == '__main__':
# 设置命令行程序
parser = argparse.ArgumentParser(description='Federated Learning')
parser.add_argument('-c', '--conf', dest='conf')
# 获取所有的参数
args = parser.parse_args()
# 读取配置文件
with open('conf.json', 'r') as f:
conf = json.load(f)
# 获取数据集, 加载描述信息
train_datasets, eval_datasets = getDatasets.get_dataset("./data", conf["type"])
# 开启服务器
server = Server(conf, eval_datasets)
# 客户端列表
clients = []
# 添加10个客户端到列表
for c in range(conf["no_models"]):
clients.append(Client(conf, server.global_model, train_datasets, c))
# 全局模型训练
for e in range(conf["global_epochs"]):
print("Global Epoch %d" % e)
# 每次训练都是从clients列表中随机采样k个进行本轮训练
candidates = random.sample(clients, conf["k"])
print("select clients is: ")
for c in candidates:
print(c.client_id)
# 权重累计
weight_accumulator = {}
# 初始化空模型参数weight_accumulator
for name, params in server.global_model.state_dict().items():
# 生成一个和参数矩阵大小相同的0矩阵
weight_accumulator[name] = torch.zeros_like(params)
# 遍历客户端，每个客户端本地训练模型
for c in candidates:
diff = c.local_train(server.global_model)
# 根据客户端的参数差值字典更新总体权重
for name, params in server.global_model.state_dict().items():
weight_accumulator[name].add_(diff[name])
# 模型参数聚合
server.model_aggregate(weight_accumulator)
# 模型评估
acc, loss = server.model_eval()
print("Epoch %d, acc: %f, loss: %f\n" % (e, acc, loss))
2.7项目结构
2.8测试
运行后的准确度以及损失如下：
三、联邦学习与中心化训练的效果对比
- 联邦训练配置：一共10台客户端设备（no_models=10），每一轮任意挑选其中的5台参与训练（k=5）， 每一次本地训练迭代次数为3次（local_epochs=3），全局迭代次数为20次（global_epochs=20）。
- 集中式训练配置：我们不需要单独编写集中式训练代码，只需要修改联邦学习配置便可使其等价于集中式训练。具体来说，我们将客户端设备no_models和每一轮挑选的参与训练设备数k都设为1即可。这样只有1台设备参与的联邦训练等价于集中式训练。其余参数配置信息与联邦学习训练一致。图
两种训练方式经过20轮迭代后的效果对比，左图是准确度对比，右图是损失函数值对比
联邦模型与单点训练模型的对比
在上图中，我们看到单点训练的模型效果（蓝条）明显要低于联邦训练的效果（绿条和橙条），这也说明了仅仅通过单个客户端的数据，不能够很好的学习到数据的全局分布特性，模型的泛化能力较差。此外，对于每一轮参与联邦训练的客户端数目（k 值）不同，其性能也会有一定的差别，k 值越大，每一轮参与训练的客户端数目越多，其性能也会越好，但每一轮的完成时间也会相对较长。