垂直联邦XGB (SecureBoost)#

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

欢迎来到SecureBoost教程!

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

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

设备设置#

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

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

[1]:
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,
    get_classic_XGB_params,
    get_classic_lightGBM_params,
)
from secretflow.ml.boost.sgb_v.model import load_model
import pprint

pp = pprint.PrettyPrinter(depth=4)

# Check the version of your SecretFlow
print('The version of SecretFlow: {}'.format(sf.__version__))
The version of SecretFlow: 1.0.0a0
[2]:
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-07-11 13:39:07,006 INFO worker.py:1544 -- Started a local Ray instance. View the dashboard at http://127.0.0.1:8265 
[3]:
alice = sf.PYU('alice')
bob = sf.PYU('bob')
heu = sf.HEU(heu_config, cluster_def['runtime_config']['field'])

数据准备#

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

[4]:
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,
)

参数准备#

[5]:
params = get_classic_XGB_params()
params['num_boost_round'] = 3
params['max_depth'] = 3
pp.pprint(params)
{'audit_paths': {},
 'base_score': 0.0,
 'batch_encoding_enabled': True,
 'bottom_rate': 0.5,
 'colsample_by_tree': 1.0,
 'early_stop_criterion_g_abs_sum': 0.0,
 'early_stop_criterion_g_abs_sum_change_ratio': 0.0,
 'enable_goss': False,
 'enable_quantization': False,
 'first_tree_with_label_holder_feature': True,
 'fixed_point_parameter': 20,
 'gamma': 0.0,
 'learning_rate': 0.3,
 'max_depth': 3,
 'max_leaf': 15,
 'num_boost_round': 3,
 'objective': 'logistic',
 'quantization_scale': 10000.0,
 'reg_lambda': 0.1,
 'rowsample_by_tree': 1.0,
 'seed': 1212,
 'sketch_eps': 0.1,
 'top_rate': 0.3,
 'tree_growing_method': 'level'}

运行 Sgb#

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

[6]:
sgb = Sgb(heu)
model = sgb.train(params, v_data, label_data)
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.factory.sgb_actor.SGBActor'> with party alice.
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.factory.sgb_actor.SGBActor'> with party bob.
(_run pid=2676570) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676570) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676570) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(_run pid=2676570) 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=2676570) 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: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.)
(_run pid=2676575) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676575) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676575) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(_run pid=2676575) 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=2676575) 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.)
(SGBActor pid=2684049) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(SGBActor pid=2684049) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(SGBActor pid=2684049) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(SGBActor pid=2684049) 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.
(SGBActor pid=2684049) 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:root:epoch 0 time 7.746685071993852s
(_run pid=2676568) [2023-07-11 13:39:22.365] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
(_run pid=2676575) [2023-07-11 13:39:22.370] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
(_run pid=2676568) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676568) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(_run pid=2676568) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(_run pid=2676568) 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=2676568) 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(heu_id=139936685360992, party=alice) pid=2682338) [2023-07-11 13:39:22.730] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
(HEUEvaluator(heu_id=139936685360992, party=bob) pid=2683835) [2023-07-11 13:39:22.755] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
INFO:root:epoch 1 time 5.602544448993285s
(_run pid=2676570) [2023-07-11 13:39:27.977] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63
INFO:root:epoch 2 time 0.3880494670011103s

模型评估#

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

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

模型保存和加载#

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

让我们先定义路径。

[8]:
# 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()
}

然后让我们保存模型。

[9]:
r = model.save_model(saving_path_dict)
wait(r)

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

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

[10]:
# 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
)

更多训练设置#

如果我们想用lightGBM的方式训练树模型怎么办?我们可以设置按叶节点训练并开启GOSS功能。

[11]:
params = get_classic_lightGBM_params()
params['num_boost_round'] = 3
params['max_leaf'] = 2**3
pp.pprint(params)
model = sgb.train(params, v_data, label_data)
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.factory.sgb_actor.SGBActor'> with party alice.
INFO:root:Create proxy actor <class 'secretflow.ml.boost.sgb_v.factory.sgb_actor.SGBActor'> with party bob.
{'audit_paths': {},
 'base_score': 0.0,
 'batch_encoding_enabled': True,
 'bottom_rate': 0.5,
 'colsample_by_tree': 1.0,
 'early_stop_criterion_g_abs_sum': 0.0,
 'early_stop_criterion_g_abs_sum_change_ratio': 0.0,
 'enable_goss': True,
 'enable_quantization': False,
 'first_tree_with_label_holder_feature': True,
 'fixed_point_parameter': 20,
 'gamma': 0.0,
 'learning_rate': 0.3,
 'max_depth': 5,
 'max_leaf': 8,
 'num_boost_round': 3,
 'objective': 'logistic',
 'quantization_scale': 10000.0,
 'reg_lambda': 0.1,
 'rowsample_by_tree': 1.0,
 'seed': 1212,
 'sketch_eps': 0.1,
 'top_rate': 0.3,
 'tree_growing_method': 'leaf'}
(SGBActor pid=2686804) INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(SGBActor pid=2686804) INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
(SGBActor pid=2686804) INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
(SGBActor pid=2686804) 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.
(SGBActor pid=2686804) 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.)
(SGBActor pid=2686804) /home/zoupeicheng.zpc/miniconda3/envs/py38/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:3652: UserWarning: 'kind' argument to argsort is ignored; only 'stable' sorts are supported.
(SGBActor pid=2686804)   warnings.warn("'kind' argument to argsort is ignored; only 'stable' sorts "
INFO:root:epoch 0 time 6.724791188011295s
INFO:root:epoch 1 time 1.0927822189987637s
INFO:root:epoch 2 time 1.002211379003711s
[12]:
yhat = model.predict(v_data)
yhat = reveal(yhat)
print(f"auc: {roc_auc_score(y, yhat)}")
auc: 0.9966901855081655

结论#

恭喜您完成了本教程!

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

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