分布式机器学习中的模型聚合

分布式机器学习中的模型聚合

我接下来的这篇论文在联邦学习(distributed learning)的上下文中引入了多任务学习，这意味着每个客户端/任务节点的训练数据分布是不同的，这样每个任务节点可以学习不同的模型，并且每个任务节点和全局模型由多个组件模型集成。本文的关键和核心在于各个任务节点所学习的模型的聚合/通信。根据模型聚合方式的不同，模型中使用的算法可以分为客户端-服务器方法和完全分散方法(实际上没有其他聚合方法，比如另一篇论文中提出的聚类方法，我们这里暂时跳过)，这两种方法都可以在具体实现中通过优化代理损失函数来替代。

论文[1]在联邦(分布式)学习的背景下引入了多任务学习，这意味着每个客户端/任务节点的训练数据分布是不同的，从而每个任务节点学习不同的模型，每个任务节点和全局模型由多个组件模型集成。本文的关键和核心在于各个任务节点所学习的模型的聚合/通信。根据模型的聚合方式不同，模型中使用的算法可以分为客户端-服务器方法和完全分散方法(实际上没有其他聚合方法，如论文[3]中提出的聚类聚合方法，代码如[4]所示，这里我们暂时跳过)。

因为要实现的任务aggregator有很多种，论文代码采取的措施(Github上开源，见[2])是先实现aggregator的抽象基类，实现一些通用的方法，并指定抽象方法的接口，然后具体的任务aggregator类继承抽象基类，再做具体的实现。

让我们首先看看任务聚合器的抽象基类。

类别聚合器(ABC):

r ' '是聚合器的基类。“聚合器”指定客户端“”之间的通信

def __init__(

自我，

客户，

global _ learners _ ensemble，

log_freq，

global_train_logger，

全局_测试_记录器，

sampling _ rate=1。

样本替换=假，

测试客户端=无，

verbose=0，

种子=无，

*参数，

**kwargs

):

rng_seed=(seed if (seed不是None，seed=0) else int(time.time()))

Self.rng=random.random (RNG种子)#随机数生成器

self . NP _ RNG=NP . random . default _ RNG(RNG _ seed)# numpy随机数生成器

如果test_clients为None:

test_clients=[]

客户端=客户端#列表[客户端]

自测客户端=测试客户端#列表[客户端]

self . global _ leaders _ ensemble=global _ leaders _ ensemble # List[leaderer]

self . device=self . global _ leaders _ ensemble . device

self.log_freq=log_freq

self.verbose=verbose

# verbose:调整输出打印的冗余度，

# ` 0 '表示安静(没有任何打印输出)，` 1 '显示日志，` 2 '显示所有本地日志；默认值为“0”。

self . global _ train _ logger=global _ train _ logger

self . global _ test _ logger=global _ test _ logger

self . model _ dim=self . global _ leaders _ aggregate . model _ dim # #模型特征维度

self.n_client

s = len(clients)
self.n_test_clients = len(test_clients)
self.n_learners = len(self.global_learners_ensemble)
# 存储为每个client分配的权重（权重为0-1之间的小数）
self.clients_weights =\
torch.tensor(
[client.n_train_samples for client in self.clients],
dtype=torch.float32
)
self.clients_weights = self.clients_weights / self.clients_weights.sum()
self.sampling_rate = sampling_rate # clients在每一轮使用的比例，默认为`1.`
self.sample_with_replacement = sample_with_replacement #对client进行采用是可重复还是无重复的，with_replacement=True表示可重复的，否则是不可重复的
# 每轮迭代需要使用到的client个数
self.n_clients_per_round = max(1, int(self.sampling_rate * self.n_clients))
# 采样得到的client列表
self.sampled_clients = list()
# 记载当前的迭代通信轮数
self.c_round = 0
self.write_logs()
@abstractmethod
def mix(self):
"""
该方法用于完成各client之间的权重参数与通信操作
"""
pass
@abstractmethod
def update_clients(self):
"""
该方法用于将所有全局分量模型拷贝到各个client，相当于boardcast操作
"""
pass
def update_test_clients(self):
"""
将全局(gobal)的所有分量模型都拷贝到各个client上
"""
def write_logs(self):
"""
对全局(global)的train和test数据集的loss和acc做记录
需要对所有client的所有样本做累加，然后除以所有client的样本总数做平均。
"""
def save_state(self, dir_path):
"""
保存aggregator的模型state,。例如, `global_learners_ensemble`中每个分量模型'learner'的state字典（以`.pt`文件格式），以及`self.clients` 中每个client的 `learners_weights` （注意，这个权重不是模型内部的参数，而是进行继承的时候对各个分量模型赋予的权重，包含train和test两部分，以一个大小为n_clients(n_test_clients)× n_learners的numpy数组的格式，即`.npy` 文件)。
"""
def load_state(self, dir_path):
"""
加载aggregator的模型state，即save_state方法里保存的那些
"""
def sample_clients(self):
"""
对clients进行采样，
如果self.sample_with_replacement为True，则为可重复采样，
否则，则为不可重复采用。
最终得到一个clients子集列表并赋予self.sampled_clients
"""

1.client-server 算法

这种方式的通信/聚合方法也称中心化(centralized)方法，因为该方法在每一轮迭代最后将所有client的权重数据汇集到server节点。这种方法的优化迭代部分的伪代码示意如下:

落实到具体代码实现上，这种方法的Aggregator设计如下：

class CentralizedAggregator(Aggregator):
    r""" 标准的中心化Aggreagator
    所有clients在每一轮迭代末和average client完全同步.
    """
    def mix(self):
        self.sample_clients()
        # 对self.sampled_clients中每个client的参数进行优化
        for client in self.sampled_clients:
            # 相当于伪代码第11行调用的LocalSolver函数
            client.step()
        # 遍历global模型(self.global_learners_ensemble) 中每一个分量模型(learner)
        # 相当于伪代码第13行
        for learner_id, learner in enumerate(self.global_learners_ensemble):
            # 获取所有client中对应learner_id的分量模型
            learners = [client.learners_ensemble[learner_id] for client in self.clients]
            # global模型的分量模型为所有client对应分量模型取平均，相当于伪代码第14行
            average_learners(learners, learner, weights=self.clients_weights)
        # 将更新后的模型赋予所有clients，相当于伪代码第5行的boardcast操作
        self.update_clients()
        # 通信轮数+1
        self.c_round += 1
        if self.c_round % self.log_freq == 0:
            self.write_logs()
    def update_clients(self):
        """
        此函数负责将所有全局分量模型拷贝到各个client，相当于伪代码中第5行的boardcast操作
        """
        for client in self.clients:
            for learner_id, learner in enumerate(client.learners_ensemble):
                copy_model(learner.model, self.global_learners_ensemble[learner_id].model)
                if callable(getattr(learner.optimizer, "set_initial_params", None)):
                    learner.optimizer.set_initial_params(
                        self.global_learners_ensemble[learner_id].model.parameters()
                    )

2. fully decentralized(完全去中心化)算法

这种方法之所以被称为去中心化的，因为该方法在每一轮迭代不需要所有client的权重数据汇集到一个特定的server节点，而只需要完成每个节点和其邻居进行通信（参数共享）即可。这种方法的优化迭代部分的伪代码示意如下:

落实到具体代码实现上，这种方法的Aggregator设计如下：

class DecentralizedAggregator(Aggregator):
    def __init__(
            self,
            clients,
            global_learners_ensemble,
            mixing_matrix,
            log_freq,
            global_train_logger,
            global_test_logger,
            sampling_rate=1.,
            sample_with_replacement=True,
            test_clients=None,
            verbose=0,
            seed=None):
        super(DecentralizedAggregator, self).__init__(
            clients=clients,
            global_learners_ensemble=global_learners_ensemble,
            log_freq=log_freq,
            global_train_logger=global_train_logger,
            global_test_logger=global_test_logger,
            sampling_rate=sampling_rate,
            sample_with_replacement=sample_with_replacement,
            test_clients=test_clients,
            verbose=verbose,
            seed=seed
        )
        self.mixing_matrix = mixing_matrix
        assert self.sampling_rate = 1, "partial sampling is not supported with DecentralizedAggregator"
    def update_clients(self):
        pass
    def mix(self):
        
        # 对各clients的模型参数进行优化
        for client in self.clients:
            client.step()
        # 存储每个模型各参数混合的权重
        # 行对应不同的client，列对应单个模型中不同的参数
        # (注意：每个分量有独立的mixing_matrix)
        mixing_matrix = torch.tensor(
            self.mixing_matrix.copy(),
            dtype=torch.float32,
            device=self.device
        )
        # 遍历global模型(self.global_learners_ensemble) 中每一个分量模型(learner)
        # 相当于伪代码第14行
        for learner_id, global_learner in enumerate(self.global_learners_ensemble):
            # 用于将指定learner_id的各client的模型state读出暂存
            state_dicts = [client.learners_ensemble[learner_id].model.state_dict() for client in self.clients]
            # 遍历global模型中的各参数
            for key, param in global_learner.model.state_dict().items():
                shape_ = param.shape
                models_params = torch.zeros(self.n_clients, int(np.prod(shape_)), device=self.device)
                for ii, sd in enumerate(state_dicts):
                    # models_params的第ii个下标存储的是第ii个client的参数
                    models_params[ii] = sd[key].view(1, -1) 
                # models_params的每一行是一个client的参数
                # @符号表示矩阵乘/矩阵向量乘
                # 故这里表示每个client参数是其他所有client参数的混合
                models_params = mixing_matrix @ models_params
                for ii, sd in enumerate(state_dicts):
                    # 将第ii个client的参数存入state_dicts中对应位置
                    sd[key] = models_params[ii].view(shape_)
            # 将更新好的参数从state_dicts存入各client节点的模型中
            for client_id, client in enumerate(self.clients):
                client.learners_ensemble[learner_id].model.load_state_dict(state_dicts[client_id])
        # 通信轮数+1
        self.c_round += 1
        if self.c_round % self.log_freq == 0:
            self.write_logs()

参考文献

• [1] Marfoq O, Neglia G, Bellet A, et al. Federated multi-task learning under a mixture of distributions[J]. Advances in Neural Information Processing Systems, 2021, 34.
• [2] https://github.com/omarfoq/FedEM
• [3] Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints
• [4] https://github.com/felisat/clustered-federated-learning

数学是符号的艺术，音乐是上界的语言。