联邦学习算法SCAFFOLD的PyTorch实现

I. 前言

SCAFFOLD的原理请见：ICML 2020 | SCAFFOLD：联邦学习的随机控制平均。

II. 数据集介绍

联邦学习中存在多个客户端，每个客户端都有自己的数据集，这个数据集他们是不愿意共享的。

数据集为某城市十个地区的风电功率，我们假设这10个地区的电力部门不愿意共享自己的数据，但是他们又想得到一个由所有数据统一训练得到的全局模型。

III. SCAFFOLD

算法伪代码：

1. 模型定义

客户端的模型为一个简单的四层神经网络模型：

# -*- coding:utf-8 -*-
"""
@Time: 2022/03/02 11:21
@Author: KI
@File: model.py
@Motto: Hungry And Humble
"""
from torch import nn


class ANN(nn.Module):
    def __init__(self, input_dim, name, B, E, lr):
        super(ANN, self).__init__()
        self.name = name
        self.B = B
        self.E = E
        self.len = 0
        self.lr = lr
        self.loss = 0
        self.fc1 = nn.Linear(input_dim, 20)
        self.control = {}
        self.delta_control = {}
        self.delta_y = {}
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout()
        self.fc2 = nn.Linear(20, 20)
        self.fc3 = nn.Linear(20, 20)
        self.fc4 = nn.Linear(20, 1)

    def forward(self, data):
        x = self.fc1(data)
        x = self.sigmoid(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        x = self.fc3(x)
        x = self.sigmoid(x)
        x = self.fc4(x)
        x = self.sigmoid(x)

        return x

2. 优化器定义

在SCAFFOLD中，本地模型的更新公式为：
$$y_i\leftarrow y_i-{\eta }_l(g_i(y_i)+c-c_i)$$
因此，优化器定义如下：

# -*- coding:utf-8 -*-
"""
@Time: 2022/03/02 13:34
@Author: KI
@File: ScaffoldOptimizer.py
@Motto: Hungry And Humble
"""
from torch.optim import Optimizer


class ScaffoldOptimizer(Optimizer):
    def __init__(self, params, lr, weight_decay):
        defaults = dict(lr=lr, weight_decay=weight_decay)
        super(ScaffoldOptimizer, self).__init__(params, defaults)

    def step(self, server_controls, client_controls, closure=None):

        for group in self.param_groups:
            for p, c, ci in zip(group['params'], server_controls.values(), client_controls.values()):
                if p.grad is None:
                    continue
                dp = p.grad.data + c.data - ci.data
                p.data = p.data - dp.data * group['lr']

核心代码为：

dp = p.grad.data + c.data - ci.data
p.data = p.data - dp.data * group['lr']

3. 服务器端

服务器端收发模型参数，并更新全局控制变量：

class Scaffold:
    def __init__(self, options):
        self.C = options['C']
        self.E = options['E']
        self.B = options['B']
        self.K = options['K']
        self.r = options['r']
        self.input_dim = options['input_dim']
        self.lr = options['lr']
        self.clients = options['clients']
        self.nn = ANN(input_dim=self.input_dim, name='server', B=self.B, E=self.E, lr=self.lr).to(
            device)
        # self.control = torch.zeros_like(self.nn.named_parameters)
        for k, v in self.nn.named_parameters():
            self.nn.control[k] = torch.zeros_like(v.data)
            self.nn.delta_control[k] = torch.zeros_like(v.data)
            self.nn.delta_y[k] = torch.zeros_like(v.data)
        self.nns = []
        for i in range(self.K):
            temp = copy.deepcopy(self.nn)
            temp.name = self.clients[i]
            temp.control = copy.deepcopy(self.nn.control)  # ci
            temp.delta_control = copy.deepcopy(self.nn.delta_control)  # ci
            temp.delta_y = copy.deepcopy(self.nn.delta_y)
            self.nns.append(temp)

    def server(self):
        for t in range(self.r):
            print('round', t + 1, ':')
            # sampling
            m = np.max([int(self.C * self.K), 1])
            index = random.sample(range(0, self.K), m)
            # dispatch
            self.dispatch(index)
            # local updating
            self.client_update(index)
            # aggregation
            self.aggregation(index)

        return self.nn

    def aggregation(self, index):
        s = 0.0
        for j in index:
            # normal
            s += self.nns[j].len
        # compute
        x = {}
        c = {}
        # init
        for k, v in self.nns[0].named_parameters():
            x[k] = torch.zeros_like(v.data)
            c[k] = torch.zeros_like(v.data)

        for j in index:
            for k, v in self.nns[j].named_parameters():
                x[k] += self.nns[j].delta_y[k] / len(index)  # averaging
                c[k] += self.nns[j].delta_control[k] / len(index)  # averaging

        # update x and c
        for k, v in self.nn.named_parameters():
            v.data += x[k].data  # lr=1
            self.nn.control[k].data += c[k].data * (len(index) / self.K)

    def dispatch(self, index):
        for j in index:
            for old_params, new_params in zip(self.nns[j].parameters(), self.nn.parameters()):
                new_params.data = old_params.data.clone()

    def client_update(self, index):  # update nn
        for k in index:
            self.nns[k] = train(self.nns[k], self.nn)

4. 客户端

客户端更新本地模型：

def train(ann, server):
    ann.train()
    Dtr, Dte = nn_seq_wind(ann.name, ann.B)
    ann.len = len(Dtr)
    print('training...')
    loss_function = nn.MSELoss().to(device)
    loss = 0
    x = copy.deepcopy(ann)
    optimizer = ScaffoldOptimizer(ann.parameters(), lr=ann.lr, weight_decay=1e-4)
    for epoch in range(ann.E):
        for (seq, label) in Dtr:
            seq = seq.to(device)
            label = label.to(device)
            y_pred = ann(seq)
            loss = loss_function(y_pred, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step(server.control, ann.control)

        print('epoch', epoch, ':', loss.item())
    # update c
    # c+ <- ci - c + 1/(E * lr) * (x-yi)
    # save ann
    temp = {}
    for k, v in ann.named_parameters():
        temp[k] = v.data.clone()

    for k, v in x.named_parameters():
        ann.control[k] = ann.control[k] - server.control[k] + (v.data - temp[k]) / (ann.E * ann.lr)
        ann.delta_y[k] = temp[k] - v.data
        ann.delta_control[k] = ann.control[k] - x.control[k]

    return ann

模型更新结束后需要更新控制变量，控制变量的更新公式为：

这里$K$为本地更新的轮数，${\eta }_l$为学习率。对应代码为：

ann.control[k] = ann.control[k] - server.control[k] + (v.data - temp[k]) / (ann.E * ann.lr)

控制变量更新结束后，计算模型和控制变量前后的差值：

ann.delta_y[k] = temp[k] - v.data
ann.delta_control[k] = ann.control[k] - x.control[k]

IV. 完整代码

完整项目我放在了GitHub上，项目地址：Scaffold-Federated Learning，原创不易，下载后请随手点下Follow和Star，感谢！！