【联邦学习实战】基于同态加密和差分隐私混合加密机制的FedAvg

前言

好久都没更新联邦学习相关内容了，这也是我更新这篇我认为非常硬核的文章的原因，这也算是实现了我在学习联邦学习半年以来的一个目标，基于混合加密机制实现联邦学习任务，这次任务使用的框架是FedAvg，在github上非常热门的联邦学习模拟实现方案，FedAvg的代码还是非常好理解的，本文的结构将主要分为三个部分，第一部分是对FedAvg代码的讲解和修改，第二部分将差分隐私机制加入到FedAvg中，包括高斯机制和拉普拉斯机制，第三部分将同态加密算法Paillier加入到第二部分中，实现基于同态加密和差分隐私混合加密机制的FedAvg，话不多说，就让我们开始吧！

---

首先把项目的代码链接附上（请为我点亮一颗Star！）

平台	链接
Github	https://github.com/heroding77/fedavg_encrypt
Gitee	https://gitee.com/heroding77/fedavg_encrypt

1. FedAvg

联邦学习有两种更新方式，对服务器端，一种是共享梯度方式，一种是共享模型参数。共享模型参数是做了几轮梯度下降，针对共享梯度，它的一大优势是通信代价会低；同时，对整个梯度信息的保护也会更好，因此本次联邦学习实战的内容也是基于共享参数形式实现的。本次实战使用的FedAvg是github上非常热门的联邦学习实现方案，它模拟了多个参与方进行联邦学习的场景，并且支持Non-IID和独立同分布数据集，代码逻辑十分清晰，对于联邦学习初学者入门再适合不过了。

FedAvg实现的是联邦学习场景下的手写数字识别模型的训练，它的结构如下图所示，包括数据集，模型，参与方和服务器端代码，以及数据集导入的代码，除了TensorFlow版本，它还包括PyTorch版本，更便于支持gpu运行代码，本文介绍的代码部分也是基于PyTorch实现的。下面就对各个部分进行简单的介绍。

1.1 getData.py

getData.py文件是对MNIST数据集的数据集提取和预处理，获得训练数据集和测试数据集，MNIST数据集的数据格式为-ubyte.gz，因此需要编写函数对图片数据进行提取，如下所示：

def extract_images(filename):
    """Extract the images into a 4D uint8 numpy array [index, y, x, depth]."""
    print('Extracting', filename)
    with gzip.open(filename) as bytestream:
        magic = _read32(bytestream)
        if magic != 2051:
            raise ValueError(
                    'Invalid magic number %d in MNIST image file: %s' %
                    (magic, filename))
        num_images = _read32(bytestream)
        rows = _read32(bytestream)
        cols = _read32(bytestream)
        buf = bytestream.read(rows * cols * num_images)
        data = np.frombuffer(buf, dtype=np.uint8)
        data = data.reshape(num_images, rows, cols, 1)
        return data

本质上是将比特流数据转换为[index, y, x, depth]维度的数组形式，index为下标，y和x分别表示每张图像行列像素点数目，depth是图像的通道，对于MNIST黑白图片，通道为1。接着编写图片标签提取函数，将标签提取出并转为one-hot形式，便于模型训练时计算损失函数并反向传播，代码如下：

# 标签one-hot编码
def dense_to_one_hot(labels_dense, num_classes=10):
    """Convert class labels from scalars to one-hot vectors."""
    num_labels = labels_dense.shape[0]
    index_offset = np.arange(num_labels) * num_classes
    labels_one_hot = np.zeros((num_labels, num_classes))
    labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
    return labels_one_hot

# 提取标签
def extract_labels(filename):
    """Extract the labels into a 1D uint8 numpy array [index]."""
    print('Extracting', filename)
    with gzip.open(filename) as bytestream:
        magic = _read32(bytestream)
        if magic != 2049:
            raise ValueError(
                    'Invalid magic number %d in MNIST label file: %s' %
                    (magic, filename))
        num_items = _read32(bytestream)
        buf = bytestream.read(num_items)
        labels = np.frombuffer(buf, dtype=np.uint8)
        return dense_to_one_hot(labels)

图片和标签信息提取完成后，多次assert以验证提取图片过程没有出现问题。接着将图片展平，即每张图片28 * 28被展平为784 * 1，这里的目的是用于全连接神经网络的使用，如果只是用卷积神经网络，大可不必将其展平（在模型部分还要复原回去）。紧接着就是对数据标准化处理，如果不进行这一步，模型训练的结果可能惨不忍睹（比如预测所有图像都是7），这是由于发生了梯度爆炸现象。标准化处理后是两个数据处理方式：IID和Non-IID，熟悉联邦学习的读者应该了解这个概念，这也是联邦学习面临的一个巨大问题——数据非独立同分布。
在传统应用场景中，数据存储在中心，ML model 可以获取所有数据的整体信息，但是在联邦学习中，由于数据仅存储在本地，导致数据之间分布的不一致性，例如，美国西海岸针叶林茂盛，东海岸阔叶林分布广泛，又因为用户之间喜好和生活习惯不同，因此产生了数据分布的不一致性; 另一方面，以家庭或亲属关系的群体之间会相互影响，产生了数据分布的不独立性，以FedAvg中MNIST为例，Non-IID的数据应该是这样表现的：参与方1只有手写数字1，参与方2只有手写数字2，以此类推。在Non-IID数据场景下，模型训练会异常缓慢，效率很低，这将在之后的实验所有介绍。
FedAvg提供了IID和Non-IID两种数据处理形式，帮助用户理解实际应用场景下联邦学习的过程，Non-IID数据就是不对数据进行操作，IID数据就是将所有数据打乱，这样各个参与方数据的分布就满足独立同分布了，代码如下：

class GetDataSet(object):
    def __init__(self, dataSetName, isIID):
        self.name = dataSetName
        self.train_data = None
        self.train_label = None
        self.train_data_size = None
        self.test_data = None
        self.test_label = None
        self.test_data_size = None

        self._index_in_train_epoch = 0

        if self.name == 'mnist':
            self.mnistDataSetConstruct(isIID)
        else:
            pass


    def mnistDataSetConstruct(self, isIID):
        data_dir = r'./data/MNIST'
        # 选定图片路径
        train_images_path = os.path.join(data_dir, 'train-images-idx3-ubyte.gz')
        train_labels_path = os.path.join(data_dir, 'train-labels-idx1-ubyte.gz')
        test_images_path = os.path.join(data_dir, 't10k-images-idx3-ubyte.gz')
        test_labels_path = os.path.join(data_dir, 't10k-labels-idx1-ubyte.gz')
        # 从.gz中提取图片
        train_images = extract_images(train_images_path)
        train_labels = extract_labels(train_labels_path)
        test_images = extract_images(test_images_path)
        test_labels = extract_labels(test_labels_path)

        assert train_images.shape[0] == train_labels.shape[0]
        assert test_images.shape[0] == test_labels.shape[0]

        self.train_data_size = train_images.shape[0]
        self.test_data_size = test_images.shape[0]

        # mnist黑白图片通道为1
        assert train_images.shape[3] == 1
        assert test_images.shape[3] == 1
        # 图片展平
        train_images = train_images.reshape(train_images.shape[0], train_images.shape[1] * train_images.shape[2])
        test_images = test_images.reshape(test_images.shape[0], test_images.shape[1] * test_images.shape[2])
        
        # 标准化处理
        train_images = train_images.astype(np.float32)
        train_images = np.multiply(train_images, 1.0 / 255.0)
        test_images = test_images.astype(np.float32)
        test_images = np.multiply(test_images, 1.0 / 255.0)

        # 是否独立同分布
        if isIID:
            # 打乱顺序
            order = np.arange(self.train_data_size)
            np.random.shuffle(order)
            self.train_data = train_images[order]
            self.train_label = train_labels[order]
        else:
            # 按照0——9顺序排列
            labels = np.argmax(train_labels, axis=1)
            order = np.argsort(labels)
            self.train_data = train_images[order]
            self.train_label = train_labels[order]

        self.test_data = test_images
        self.test_label = test_labels

1.2 Models.py

在Models.py文件中，定义了执行任务所使用的模型，这里提供了两种模型，简单全连接层模型和简单卷积神经网络模型，每个模型都继承了torch的nn.Module，表明是神经网络模型，每个模型类中都定义了__init__(self)和forward()两个函数，初始化函数中定义了每个模型的组件，比如简单全连接层模型中定义了三个卷积模块，两个池化模块，两个线性模块，并在forward()函数中组装起来。代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

'''
简单全连接模型
'''
class Mnist_2NN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 200)
        self.fc2 = nn.Linear(200, 200)
        self.fc3 = nn.Linear(200, 10)

    def forward(self, inputs):
        tensor = F.relu(self.fc1(inputs))
        tensor = F.relu(self.fc2(tensor))
        tensor = self.fc3(tensor)
        return tensor

'''
简单卷积神经网络模型
'''
class Mnist_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        # 定义每一层模型
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=0)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=0)
        self.fc1 = nn.Linear(3*3*64, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, inputs):
        # 构造模型
        tensor = inputs.view(-1, 1, 28, 28)
        tensor = F.relu(self.conv1(tensor))
        tensor = self.pool1(tensor)
        tensor = F.relu(self.conv2(tensor