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 ← y i ? η l ( g i ( y i ) + c ? c i ) y_i \gets y_i-\eta_l(g_i(y_i)+c-c_i) yi?←yi??ηl?(gi?(yi?)+c?ci?)
因此，優化器定義如下：

# -*- 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 K K為本地更新的輪數， η l \eta_l η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，感謝！！