联邦学习实战-1:用python从零开始实现横向联邦学习

什么是联邦学习？

简单来说就是在一个多方的环境中，数据集是零散的（在各个不同的客户端中），那么怎样实现机器学习算法呢？

首先想到的就是将多个数据集合并合并起来，然后统一的使用传统的机器学习或者深度学习算法进行计算，但是如果有一方因为数据隐私问题不愿意提交自己的数据呢？

那么就出现了联邦学习，核心就是“数据不动模型动，数据可用不可见”

多个客户端不提交数据而是提交训练时的参数/梯度给中心服务器，中心服务器进行计算后再将参数/梯度返回多个客户端再学习的过程

整个过程数据的所有权依然在用户手中，这就是联邦学习

当然数据隐私方面，联邦学习还将结合同态加密、安全多方计算、查分隐私等隐私计算技术实现更安全的保障

（ps：这里只是简单的介绍，详细的内容请多查阅其他资料）

基本概念入门学习见：《Federated_Machine_Learning:Concept_and_Applications》精读

一、环境准备

实验基于机器学习库PyTorch, 所以需要一些基础的PyTorch使用

（ps：不会也没事，下面代码有详细的注释，因为我也刚刚入门 ）

• anaconda、python3.7、PyTorch
  pip install torch
• GPU安装CUDA、cuDNN

二、横向联邦图像分类

基本信息

数据集：CIFAR10

模型：ResNet-18

环境角色:

• 中心服务器
• 多个客户端

为了简化，这里服务器客户端都是在单机上模拟，后面使用FATE会在真实多台机器上实现

基本的流程：

1. 服务器按配置生成初始化模型，客户端按照顺序将数据集横向不重叠切割
2. 服务器将全局模型发送给客户端
3. 客户端接收全局模型（来自服务器）通过本地多次迭代计算本地参数差值返回给服务器
4. 服务器聚合各个客户端差值更新模型，再评估当前模型性能
5. 如果性能未达标，则重复2过程，否则结束

2.1 配置文件

配置文件包含了整个项目的模型、数据集、epoch等核心训练参数

需要注意的是，一般来说配置文件需要在所有的客户端与服务端之间同步一致

创建一个配置文件:

项目文件夹下./utils/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 
}

• model_name：模型名称
• no_models：客户端总数量
• type：数据集信息
• global_epochs：全局迭代次数，即服务端与客户端的通信迭代次数
• local_epochs：本地模型训练迭代次数
• k：每一轮迭代时，服务端会从所有客户端中挑选k个客户端参与训练。
• batch_size：本地训练每一轮的样本数
• lr，momentum，lambda：本地训练的超参数设置

2.1 构建训练数据集

构建数据集代码如下:

datasets.py

import torchvision as tv

# 获取数据集
def get_dataset(dir, name):
    if name == 'mnist':
        # root: 数据路径
        # train参数表示是否是训练集或者测试集
        # download=true表示从互联网上下载数据集并把数据集放在root路径中
        # transform：图像类型的转换
        train_dataset = tv.datasets.MNIST(dir, train=True, download=True, transform=tv.transforms.ToTensor())
        eval_dataset = tv.datasets.MNIST(dir, train=False, transform=tv.transforms.ToTensor())
    elif name == 'cifar':
        # 设置两个转换格式
        # transforms.Compose 是将多个transform组合起来使用（由transform构成的列表）
        transform_train = tv.transforms.Compose([
            # transforms.RandomCrop： 切割中心点的位置随机选取
            tv.transforms.RandomCrop(32, padding=4), tv.transforms.RandomHorizontalFlip(),
            tv.transforms.ToTensor(),
            # transforms.Normalize： 给定均值：(R,G,B) 方差：（R，G，B），将会把Tensor正则化
            tv.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
        transform_test = tv.transforms.Compose([
            tv.transforms.ToTensor(),
            tv.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
        train_dataset = tv.datasets.CIFAR10(dir, train=True, download=True, transform=transform_train)
        eval_dataset = tv.datasets.CIFAR10(dir, train=False, transform=transform_test)
    return train_dataset, eval_dataset

2.2 服务端

服务端的主要功能是模型的聚合、评估，最终的模型也是在服务器上生成

首先创建一个服务类

所有的程序放在server.py

构造函数

定义其构造函数：

# 定义构造函数
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,
    # 设置单个批次大小32
    batch_size=self.conf["batch_size"],
    # 打乱数据集
    shuffle=True
  )

聚合函数

定义全局联邦平均FedAvg聚合函数：

FedAvg算法的公式如下：

$G^{t+1}=G^t+\lambda {\sum }_{i=1}^m(L_i^{t+1}-G_i^t)$​

$G^t$表示第t轮更新的全局模型参数，$L_i^{t+1}$​表示第i个客户端在第t+1轮本地更新后的模型

在模型聚合时，weight_accumulator就是$(L_i^{t+1}-G_i^t)i=1,2,...m$​部分，具体weight_accumulator的计算会在后面详细介绍其实现

# 全局聚合模型
# weight_accumulator 存储了每一个客户端的上传参数变化值/差值
def model_aggregate(self, weight_accumulator):
  # 遍历服务器的全局模型
  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]
            # 统计预测结果与真实标签target的匹配总个数
            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.3 客户端

客户端的主要功能是：

• 接受服务器下发的指令和全局模型
• 利用本地数据进行局部模型训练

此部分所有程序都在client.py中

构造函数

定义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每个batch加载多少个样本(默认: 1)
            batch_size=conf["batch_size"],
            # 指定子集合
            # sampler定义从数据集中提取样本的策略
            sampler=torch.utils.data.sampler.SubsetRandomSampler(indices)
        )

本案例中根据ID将数据集进行横向切分，每个客户端之间没有交集

本地训练

本地模型训练函数：采用交叉熵作为本地训练的损失函数，并使用梯度下降来求解参数

		# 模型本地训练函数
    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()
                # 梯度
                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.4 整合

所有程序代码在main.py中

import argparse
import json
import random

import datasets
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(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))

    print("\n\n")

    # 全局模型训练
    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.5 测试

按照以上配置，(本人)运行后的准确度以及损失为：

官方的对比：

联邦学习与中心化训练的效果对比

• 联邦训练配置：一共10台客户端设备（no_models=10），每一轮任意挑选其中的5台参与训练（k=5）， 每一次本地训练迭代次数为3次（local_epochs=3），全局迭代次数为20次（global_epochs=20）。
• 集中式训练配置：我们不需要单独编写集中式训练代码，只需要修改联邦学习配置既可使其等价于集中式训练。具体来说，我们将客户端设备no_models和每一轮挑选的参与训练设备数k都设为1即可。这样只有1台设备参与的联邦训练等价于集中式训练。其余参数配置信息与联邦学习训练一致。图中我们将局部迭代次数分别设置了1，2，3来进行比较。

联邦学习在模型推断上的效果对比

图中的单点训练只的是在某一个客户端下，利用本地的数据进行模型训练的结果。

• 我们看到单点训练的模型效果（蓝色条）明显要低于联邦训练 的效果（绿色条和红色条），这也说明了仅仅通过单个客户端的数据，不能够 很好的学习到数据的全局分布特性，模型的泛化能力较差。
• 此外，对于每一轮 参与联邦训练的客户端数目（k 值）不同，其性能也会有一定的差别，k 值越大，每一轮参与训练的客户端数目越多，其性能也会越好，但每一轮的完成时间也会相对较长。

学习资料来自于：

杨强：《联邦学习实战》

https://github.com/FederatedAI/Practicing-Federated-Learning/tree/main/chapter03_Python_image_classification