垂直联邦XGB (SecureBoost)

以下代码仅供演示。出于系统安全考虑，请 不要 直接用于生产。

欢迎来到SecureBoost教程！

在本教程中，我们将探索如何使用隐语的树模型能力，使用SecureBoost算法执行垂直联邦学习。SecureBoost是一种经典算法，它优先保护垂直分区数据集中的标签信息。它使用同态加密技术实现标签加密和密文中的关键树增强步骤执行。其结果是由PYU对象组成的分布式提升树模型，每个参与方仅了解自己的拆分点。该实现利用HEU和PYU设备实现高性能。

让我们深入了解细节，学习如何使用隐语进行SecureBoost！

设备设置

与其他算法类似，设置安全集群和指定设备对于SecureBoost的实现是必要的。

特别是，必须指定一个HEU设备以确保SecureBoost中标签的加密和敏感信息的保护。

import spu
from sklearn.metrics import roc_auc_score

import secretflow as sf
from secretflow.data import FedNdarray, PartitionWay
from secretflow.device.driver import reveal, wait
from secretflow.ml.boost.sgb_v import Sgb
from secretflow.ml.boost.sgb_v.model import load_model

alice_ip = '127.0.0.1'
bob_ip = '127.0.0.1'
ip_party_map = {bob_ip: 'bob', alice_ip: 'alice'}

_system_config = {'lineage_pinning_enabled': False}
sf.shutdown()
# init cluster
sf.init(
    ['alice', 'bob'],
    address='local',
    _system_config=_system_config,
    object_store_memory=5 * 1024 * 1024 * 1024,
)

# SPU settings
cluster_def = {
    'nodes': [
        {'party': 'alice', 'id': 'local:0', 'address': alice_ip + ':12945'},
        {'party': 'bob', 'id': 'local:1', 'address': bob_ip + ':12946'},
        # {'party': 'carol', 'id': 'local:2', 'address': '127.0.0.1:12347'},
    ],
    'runtime_config': {
        # SEMI2K support 2/3 PC, ABY3 only support 3PC, CHEETAH only support 2PC.
        # pls pay attention to size of nodes above. nodes size need match to PC setting.
        'protocol': spu.spu_pb2.SEMI2K,
        'field': spu.spu_pb2.FM128,
    },
}

# HEU settings
heu_config = {
    'sk_keeper': {'party': 'alice'},
    'evaluators': [{'party': 'bob'}],
    'mode': 'PHEU',
    'he_parameters': {
        # ou is a fast encryption schema that is as secure as paillier.
        'schema': 'ou',
        'key_pair': {
            'generate': {
                # bit size should be 2048 to provide sufficient security.
                'bit_size': 2048,
            },
        },
    },
    'encoding': {
        'cleartext_type': 'DT_I32',
        'encoder': "IntegerEncoder",
        'encoder_args': {"scale": 1},
    },
}

2023-04-19 13:56:51,630 INFO worker.py:1544 -- Started a local Ray instance. View the dashboard at http://127.0.0.1:8265 

alice = sf.PYU('alice')
bob = sf.PYU('bob')
heu = sf.HEU(heu_config, cluster_def['runtime_config']['field'])

(raylet) /home/zoupeicheng.zpc/miniconda3/envs/py38/lib/python3.8/site-packages/ray/dashboard/modules/reporter/reporter_agent.py:56: UserWarning: `gpustat` package is not installed. GPU monitoring is not available. To have full functionality of the dashboard please install `pip install ray[default]`.)
(raylet)   warnings.warn(

数据准备

我们将准备一个垂直数据集。

from sklearn.datasets import load_breast_cancer

ds = load_breast_cancer()
x, y = ds['data'], ds['target']

v_data = FedNdarray(
    {
        alice: (alice(lambda: x[:, :15])()),
        bob: (bob(lambda: x[:, 15:])()),
    },
    partition_way=PartitionWay.VERTICAL,
)
label_data = FedNdarray(
    {alice: (alice(lambda: y)())},
    partition_way=PartitionWay.VERTICAL,
)

参数准备

params = {
    'num_boost_round': 5,
    'max_depth': 5,
    # about 13 bin numbers
    'sketch_eps': 0.08,
    # use 'linear' if want to do regression
    # for classification, currently only supports binary classfication
    'objective': 'logistic',
    'reg_lambda': 0.3,
    'subsample': 0.9,
    'colsample_by_tree': 0.9,
    # pre-pruning parameter. splits with gain value less than it will be pruned.
    'gamma': 1,
}

运行 Sgb

我们使用 heu 设备创建一个 Sgb 对象，并拟合数据。

sgb = Sgb(heu)
model = sgb.train(params, v_data, label_data)

INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.core.split_tree_trainer.split_tree_trainer.SplitTreeTrainer'> with party alice.
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.core.split_tree_trainer.split_tree_trainer.SplitTreeTrainer'> with party bob.
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.core.label_holder.label_holder.LabelHolder'> with party alice.
INFO:root:global_setup time 1.6802808750071563s
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Remote TPU is not linked into jax; skipping remote TPU.
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(LabelHolder pid=3311307) INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
(LabelHolder pid=3311307) WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

(HEUSkKeeper pid=3309706) [2023-04-19 13:57:02.080] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
(HEUEvaluator pid=3311083) [2023-04-19 13:57:02.127] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63

INFO:root:epoch 0 time 3.0825682659633458s

(_run pid=3303716) [2023-04-19 13:57:02.636] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
(_run pid=3303717) [2023-04-19 13:57:02.630] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63

INFO:root:epoch 1 time 0.5092060890165158s
INFO:root:epoch 2 time 0.4486289119813591s
INFO:root:epoch 3 time 0.47898129001259804s
INFO:root:epoch 4 time 0.3945654829731211s

模型评估

现在我们可以将模型输出与真实标签进行比较。

yhat = model.predict(v_data)
yhat = reveal(yhat)
print(f"auc: {roc_auc_score(y, yhat)}")

(_run pid=3303717) INFO:jax._src.xla_bridge:Remote TPU is not linked into jax; skipping remote TPU.
(_run pid=3303717) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
(_run pid=3303717) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=3303717) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=3303717) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(_run pid=3303717) INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
(_run pid=3303717) WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
INFO:jax._src.xla_bridge:Remote TPU is not linked into jax; skipping remote TPU.
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
INFO:jax._src.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

auc: 0.9994979123724961

模型保存和加载

我们现在可以保存模型, 并在以后使用它。请注意，模型是分布式的，我们将保存到多个参与方，并从多个参与方中加载。

让我们先定义路径。

# each participant party needs a location to store
saving_path_dict = {
    # in production we may use remote oss, for example.
    device: "./" + device.party
    for device in v_data.partitions.keys()
}

然后让我们保存模型。

r = model.save_model(saving_path_dict)
wait(r)

现在您可以在之前指定的位置检查文件。

最后，让我们加载模型并进行一次检查。

# alice is our label holder
model_loaded = load_model(saving_path_dict, alice)
fed_yhat_loaded = model_loaded.predict(v_data, alice)
yhat_loaded = reveal(fed_yhat_loaded.partitions[alice])

assert (
    yhat == yhat_loaded
).all(), "loaded model predictions should match original, yhat {} vs yhat_loaded {}".format(
    yhat, yhat_loaded
)

结论

恭喜您完成了本教程！

在本教程中，我们学习了如何在隐语中使用树模型进行训练，并探索了 SecureBoost，这是一种专门为垂直分区数据集设计的高性能提升算法。SecureBoost 类似于 XGBoost，但重点关注在垂直学习场景中保护敏感标签。通过利用同态加密和 PYUObjects，SecureBoost 允许我们训练强大的分布式森林模型，同时保护数据的隐私和安全。

感谢您参与本教程，希望您觉得它充满启发和帮助！