

只看这三点就够：一文了解联邦学习框架

京东科技集团研究院

2020-11-19

导读：联邦学习其核心就是一个分布式的机器学习。

来源 / 京东数科风险算法与技术

近年来，联邦学习(Federated Learning)井喷式的概念性文章层出不穷，当联邦学习的理念已经逐渐被行业认可到加以重视，从理论走向实践便是一条需要快马加鞭之路。

联邦学习其核心就是一个分布式的机器学习。通过传参数，不上传数据的方式做分布式的机器学习，相较于传统分布式机器学习，其实现了数据隐私保护。通过整合各个节点上的参数, 不同的设备可以在保持设备中大部分数据的同时，实现模型训练更新。当前市场上已经出现了一些联邦学习框架，但能真正用于实际生产的屈指可数。正是因为联邦学习涉及的技术领域之多，并且需要能够兼顾系统性能、使用资源等重要指标，所以在实际落地过程中对技术人员专业素质要求很高。

本文从部署、框架角度，总结了联邦学习工程落地与实现的部分内容，目的是帮助入门者快速了解联邦学习工程大致框架。

1、部署

在实际落地时，不同企业间往往无法达到机器配置和网络环境的统一，因此会造成较高的先期训练及推理环境准备带来的的人力成本，从以往的经验来看，较为便捷的方式是通过打包联邦学习应用以及依赖包，到一个可移植的轻量容器——docker中，让其可以在任何安装了Docker的机器上运行，而不用关心底层操作系统，这类似船舶使用的集装箱，可以便捷地装配到轮船、汽车等运载工具上。

这种做法往往具有以下优点：

不仅能节约时间，快速部署和启动（秒级甚至毫秒级），还能节约成本，相于较虚拟机动辄几个G的磁盘空间，docker容器可以减少到MB级；

方便部署，直接运行已经配好的容器，解决开发人员由于安装环境带来的部署困难；docker的镜像提供了除内核外完整的运行时环境，确保环境一致性，从而不会在出现“这段代码在我机器上没问题”这类问题；

方便持续集成，通过与代码进行关联使持续集成非常方便；

方便构建基于SOA架构或微服务架构的系统，通过服务编排，更好的松耦合；

标准化应用发布，可以多平台部署。

2、训练

在实际生产中进行联邦学习训练时，不仅要考虑到系统的耦合性、稳定性，还要考虑业务需求对接多个数据源，需要支持市面上的各种计算引擎如Spark、Flink、Storm等，抑或需要满足高可用进行集群部署等等。因此，为了应对复杂的业务需求，对系统的各个组件，我们往往需要在灵活与便捷之间寻找一个平衡点。

模型训练往往较为耗时，这对系统的稳定性要求极高，一个完备的训练服务应该支持前期校验、中期容错、后期重试、全周期监控、信息可追踪等功能特性。工程落地涉及众多问题，面临大量挑战，真正能够实现落地便已非常复杂，实际生产应用时出现的问题更待陆续解决。如在当前的联邦学习实践场景下，训练过程中的“断点续跑”便是一个重要待解决难题。在实际业务场景中，一个模型的训练往往长达数十个小时，在此过程中，如果由于网络不稳定等问题导致训练过程中断，则该模型需要重新训练，这将极大地增加训练模型的时间成本，带来十分不友好的用户体验。

一个完整的训练框架常需要支持以下服务：

▲图一训练框架示意图

通信服务：由于端与端之间需要通信，为了能尽可能少地向对方暴露我方服务的信息，以及简便性调用训练服务，我们需要引入网关服务实现服务路由，对外暴露gRPC接口以及HTTP接口，外部系统的所有请求都将委托给网关服务进行请求转发。

训练服务：该服务包括校验组件、任务调度组件，元数据中心管理组件，联邦学习组件。校验组件：负责校验我们提交的配置参数。任务调度组件：负责解析配置参数，以及进行整个训练任务的调度。在这里我们引入了设计模式中的责任链模式，按照我们指定的组件运行顺序，将一个训练任务转化成一条责任链，并提交给任务线程池去执行。元数据中心管理组件：负责记录每个训练任务的进度以及运行状态，配置参数，以及联邦学习合作方以及角色。联邦学习组件：这里面有很多小组件，用来完成模型训练过程中需要的各种功能。

模型管理服务：当完成训练任务后，训练服务将训练好的模型信息发送给模型管理服务，由模型管理服务完成持久化存储，分组等操作。

注册中心：为了保证服务的高可用，引入如ZooKeeper的注册服务，当服务启动的时候，会将服务信息注册到ZooKeeper中，当向网关服务发起训练请求的时候，网关服务会从ZooKeeper中拉取到可用的服务，通过指定的负载均衡策略完成服务调用。

在此基础上，训练流程大致可以分为如下几个步骤：

▲图二训练流程示意图

提交训练任务：当我们向网关服务提交一个训练任务之后，网关服务会将请求路由到训练服务。

检验参数：在训练服务接受到网关服务请求之后，开始训练之前，我们会先检验配置参数的准确性，比如说需要的算法组件库是否存在、训练的参数是否合理、格式是否正确等，当校验通过后，训练服务会解析我们上传的配置参数，根据我们的配置去实例化所需要的联邦学习组件。

加载样本数据：将不同类型的数据源（如: CSV文件、HDFS、MySQL）的样本数据转换为训练服务可识别的格式类型。

特征数据相交：比如A方的用户ID为u1，u2，u3，u4，而B方的用户名为u1，u2，u3，u5。交集后，A方和B方知道他们相同的用户ID，分别为u1，u2，u3，但A方对B方的其他用户ID（如u5）一无所知，B方对A方其它用户（u4）一无所知。

联邦训练：指定联邦算法组件（如：LR、DNN）进行联合模型训练。

模型评估：评估组件提供一些用于分类和回归的评估方法，它包含AUC、KS等指标用来评估模型。

模型存储：完成训练任务后，训练服务会将训练好的模型信息发送给模型管理服务，由模型管理服务将模型进行分类，版本管理等等，然后将模型信息存储到存储服务中，并保存模型的元数据信息、地址等。

最后当我们将模型持久化存储后，我们可以通过评估组件选择出效果最好的模型，然后将模型导入到推理环境中进行预测，这样我们就可以将联邦学习的成果运用到了企业生产了。

3、推理

在我们完成联邦建模任务并且成功将最终训练好的模型存储在相应的存储模块后，发起方A端就可以发起推理任务，在与B端的配合下展开预测，根据双方的特征值，得到预测结果，而这个结果正是我们联合建模的最终目的。

与训练不同，由于具有实时调用场景，往往对推理服务的性能有很高的要求，并且要求推理服务能在不影响线上实时调用的基础上，实现无感知的服务升级。企业间的推理预测，往往涉及对账问题，于是对端监控变得极为重要，这能保证在推理完成后，合作方之间能够根据监控信息做到有帐可查，并能排查推理失败等异常情况。

一个完整的推理框架通常需要支持以下服务：

▲图三推理框架示意图

通信服务：端与端之间需要通信，所以我们需要一个代理，对外暴露gRPC接口跟HTTP接口，路由的转发以及外部系统的所有请求都委托给这一个代理，同时也可以根据业务的特点，决定负载均衡的方式，比如在接入代理之前部署Nginx来实现反向代理。

推理服务：当一个新的推理请求过来时，它可以将所有推理相关的接口注册到如ZooKeeper的注册中心，外部系统可以通过服务发现获取接口地址加以调用；从远程存储系统中存储模型到本地，根据请求信息从本地存储系统中选择模型并加载模型；匹配到需要参与预测的用户信息，开展预测任务。

模型管理服务：完成模型持久化存储，分组等，保存训练好的模型信息。

注册中心：为了保证服务的高可用，引入如ZooKeeper的注册服务，当服务启动的时候，会将服务信息注册，当向网关服务发起推理请求的时候，网关服务会从注册服务中拉取到可用的服务，通过指定的负载均衡策略完成服务调用。

存储服务：我们还需要将每次预测的结果存储起来以满足一些业务的需求，同时也需要将模型存储起来，持久化模型到本地可以保证在推理中某些组件发生问题时快速恢复过来，以及不需要每次发起推理请求时都从分布式存储系统中加载模型，从而保证安全也提高了效率。

在此基础上，整个推理的流程大致可以分为以下几个步骤：

▲图四推理流程示意图

提交训练任务：当我们向网关服务提交一个推理任务之后，网关服务会将请求路由到推理服务。

获取模型：如果是第一次加载该模型，则会将其从分布式存储系统持久化到本地的存储系统系统中，如果已经被加载过，则直接从本地存储系统中读取。

预处理&预测：在A，B都加载完模型后，B端根据Featureid从本侧外部系统中获取相应的特征信息，双方都需要对本地特征进行预处理，然后开始预测，并将结果通过通信服务传给A端。

后期处理：A端将B端的结果进行整合，在本地进行规则映射等后期处理，并将最终的结果保存。

综上所述，整个推理流程都比较简单，但是当我们以工程的角度去看的时候是远远不够的，不管考虑到系统的耦合性还是稳定性，例如，当业务需求是要将整个推理从联邦系统中抽出来单独部署的时候，当我们要查询不同模型不同版本预测的历史结果的时候，又或者需要满足高可用进行集群部署的时候等等。因此为了应对复杂的业务需求，对系统的各个组件，我们亦需在灵活与便捷之间寻找一个平衡点。由于在实际生产中有着复杂的业务需求以及不可预测的问题，因此需要我们根据实际情况完善架构，例如添加相应的服务治理功能，以及线程池的规划等等。