def traverse_a_tree(tree: HeteroDecisionTreeHost, sample, cur_node_idx):
nid, _ = tree.traverse_tree(predict_state=(cur_node_idx, -1), data_inst=sample,
decoder=tree.decode, split_maskdict=tree.split_maskdict,
missing_dir_maskdict=tree.missing_dir_maskdict, sitename=tree.sitename,
tree_=tree.tree_node, zero_as_missing=tree.zero_as_missing,
use_missing=tree.use_missing)
return nid, _
def load_hetero_tree_learner(role,
tree_param,
model_meta,
model_param,
flow_id,
runtime_idx,
host_party_list=None,
fast_sbt=False,
tree_type=None,
target_host_id=None):
if role == consts.HOST:
if fast_sbt:
tree = HeteroFastDecisionTreeHost(tree_param)
else:
tree = HeteroDecisionTreeHost(tree_param)
tree.load_model(model_meta, model_param)
tree.set_flowid(flow_id)
tree.set_runtime_idx(runtime_idx)
if fast_sbt:
tree.set_tree_work_mode(tree_type, target_host_id)
tree.set_self_host_id(runtime_idx)
elif role == consts.GUEST:
if fast_sbt:
tree = HeteroFastDecisionTreeGuest(tree_param)
else:
tree = HeteroDecisionTreeGuest(tree_param)
tree.load_model(model_meta, model_param)
tree.set_flowid(flow_id)
tree.set_runtime_idx(runtime_idx)
tree.set_host_party_idlist(host_party_list)
if fast_sbt:
tree.set_tree_work_mode(tree_type, target_host_id)
else:
raise ValueError('unknown role: {}'.format(role))
return tree
def fit_a_booster(self, epoch_idx: int, booster_dim: int):
tree = HeteroDecisionTreeHost(tree_param=self.tree_param)
tree.init(flowid=self.generate_flowid(epoch_idx, booster_dim),
valid_features=self.sample_valid_features(),
data_bin=self.data_bin, bin_split_points=self.bin_split_points,
bin_sparse_points=self.bin_sparse_points,
runtime_idx=self.component_properties.local_partyid,
goss_subsample=self.enable_goss,
bin_num=self.bin_num,
complete_secure=True if (self.complete_secure and epoch_idx == 0) else False,
cipher_compressing=self.cipher_compressing,
new_ver=self.new_ver
)
tree.fit()
return tree
def load_booster(self, model_meta, model_param, epoch_idx, booster_idx):
tree = HeteroDecisionTreeHost(self.tree_param)
tree.load_model(model_meta, model_param)
tree.set_flowid(self.generate_flowid(epoch_idx, booster_idx))
tree.set_runtime_idx(self.component_properties.local_partyid)
return tree
def fit_a_booster(self, epoch_idx: int, booster_dim: int):
self.check_run_sp_opt()
tree = HeteroDecisionTreeHost(tree_param=self.tree_param)
tree.set_input_data(data_bin=self.data_bin,
bin_split_points=self.bin_split_points,
bin_sparse_points=self.bin_sparse_points)
tree.set_valid_features(self.sample_valid_features())
tree.set_flowid(self.generate_flowid(epoch_idx, booster_dim))
tree.set_runtime_idx(self.component_properties.local_partyid)
if self.run_sparse_opt:
tree.activate_sparse_hist_opt()
tree.set_dense_data_for_sparse_opt(
data_bin_dense=self.data_bin_dense, bin_num=self.bin_num)
if self.complete_secure and epoch_idx == 0:
tree.set_as_complete_secure_tree()
tree.fit()
return tree
def produce_hetero_tree_learner(
role,
tree_param: DecisionTreeParam,
flow_id,
data_bin,
bin_split_points,
bin_sparse_points,
task_type,
valid_features,
host_party_list,
runtime_idx,
cipher_compress=True,
mo_tree=False,
class_num=1,
g_h=None,
encrypter=None, # guest only
goss_subsample=False,
complete_secure=False,
max_sample_weights=1.0,
bin_num=None, # host only
fast_sbt=False,
tree_type=None,
target_host_id=None, # fast sbt only
guest_depth=2,
host_depth=3 # fast sbt only
):
if role == consts.GUEST:
if not fast_sbt:
tree = HeteroDecisionTreeGuest(tree_param)
else:
tree = HeteroFastDecisionTreeGuest(tree_param)
tree.set_tree_work_mode(tree_type, target_host_id)
tree.set_layered_depth(guest_depth, host_depth)
tree.init(flowid=flow_id,
data_bin=data_bin,
bin_split_points=bin_split_points,
bin_sparse_points=bin_sparse_points,
grad_and_hess=g_h,
encrypter=encrypter,
task_type=task_type,
valid_features=valid_features,
host_party_list=host_party_list,
runtime_idx=runtime_idx,
goss_subsample=goss_subsample,
complete_secure=complete_secure,
cipher_compressing=cipher_compress,
max_sample_weight=max_sample_weights,
mo_tree=mo_tree,
class_num=class_num)
elif role == consts.HOST:
if not fast_sbt:
tree = HeteroDecisionTreeHost(tree_param)
else:
tree = HeteroFastDecisionTreeHost(tree_param)
tree.set_tree_work_mode(tree_type, target_host_id)
tree.set_layered_depth(guest_depth, host_depth)
tree.set_self_host_id(runtime_idx)
tree.set_host_party_idlist(host_party_list)
tree.init(flowid=flow_id,
valid_features=valid_features,
data_bin=data_bin,
bin_split_points=bin_split_points,
bin_sparse_points=bin_sparse_points,
runtime_idx=runtime_idx,
goss_subsample=goss_subsample,
complete_secure=complete_secure,
cipher_compressing=cipher_compress,
bin_num=bin_num,
mo_tree=mo_tree)
else:
raise ValueError('unknown role: {}'.format(role))
return tree