def train(args):
if args.verbosity == 0:
logging.basicConfig(level=logging.WARNING)
elif args.verbosity == 1:
logging.basicConfig(level=logging.INFO)
else:
logging.basicConfig(level=logging.DEBUG)
assert args.role in ['leader', 'follower', 'local'], \
"role must be leader, follower, or local"
assert args.mode in ['train', 'test', 'eval'], \
"mode must be train, test, or eval"
if args.data_path.endswith('.csv'):
with open(args.data_path, 'rb') as fin:
data = np.loadtxt(fin, delimiter=',')
if args.mode == 'train' or args.mode == 'test':
if args.role == 'leader' or args.role == 'local':
X = data[:, :-1]
y = data[:, -1]
else:
X = data
y = None
else: # eval
X = data
y = None
else:
raise ValueError("Unsupported data type %s" % args.data_path)
if args.role != 'local':
bridge = Bridge(args.role, int(args.local_addr.split(':')[1]),
args.peer_addr, args.application_id, 0)
else:
bridge = None
booster = BoostingTreeEnsamble(bridge,
learning_rate=args.learning_rate,
max_iters=args.max_iters,
max_depth=args.max_depth,
l2_regularization=args.l2_regularization,
max_bins=args.max_bins,
num_parallel=args.num_parallel)
if args.load_model_path:
booster.load_saved_model(args.load_model_path)
if args.mode == 'train':
booster.fit(X, y, args.checkpoint_path)
elif args.mode == 'test':
pred = booster.batch_predict(X)
acc = sum((pred > 0.5) == y) / len(y)
logging.info("Test accuracy: %f", acc)
else:
pred = booster.batch_predict(X)
for i in pred:
print(i)
if args.export_path:
booster.save_model(args.export_path)
def local_test_boosting_tree_helper(self, X, y):
booster = BoostingTreeEnsamble(None,
max_iters=5,
max_depth=3,
num_parallel=2)
booster.fit(X, y)
pred = booster.batch_predict(X)
return pred
def follower_test_boosting_tree_helper(self, X, cat_X):
bridge = fl.trainer.bridge.Bridge('follower',
50052,
'localhost:50051',
streaming_mode=False)
booster = BoostingTreeEnsamble(bridge, max_iters=3, max_depth=2)
booster.fit(X, None, cat_features=cat_X)
pred = booster.batch_predict(X, cat_features=cat_X)
bridge.terminate()
def local_test_boosting_tree_helper(self, X, y, cat_X):
booster = BoostingTreeEnsamble(None,
max_iters=3,
max_depth=2,
num_parallel=2)
train_pred = booster.fit(X, y, cat_features=cat_X)
pred = booster.batch_predict(X, cat_features=cat_X)
np.testing.assert_almost_equal(train_pred, pred)
return pred
def follower_test_boosting_tree_helper(self, X, cat_X):
bridge = fl.trainer.bridge.Bridge('follower',
50052,
'localhost:50051',
streaming_mode=True)
booster = BoostingTreeEnsamble(bridge, max_iters=3, max_depth=2)
booster.fit(X, None, cat_features=cat_X)
pred = booster.batch_predict(X, cat_features=cat_X, get_raw_score=True)
bridge.terminate()
np.testing.assert_almost_equal(pred, 0)
def leader_test_boosting_tree_helper(self, X, y):
bridge = fl.trainer.bridge.Bridge('leader',
50051,
'localhost:50052',
streaming_mode=False)
booster = BoostingTreeEnsamble(bridge, max_iters=5, max_depth=3)
booster.fit(X, y)
pred = booster.batch_predict(X)
bridge.terminate()
return pred
def leader_test_boosting_tree_helper(self, X, y, cat_X):
bridge = fl.trainer.bridge.Bridge('leader',
50051,
'localhost:50052',
streaming_mode=False)
booster = BoostingTreeEnsamble(bridge, max_iters=3, max_depth=2)
train_pred = booster.fit(X, y, cat_features=cat_X)
pred = booster.batch_predict(X, cat_features=cat_X)
np.testing.assert_almost_equal(train_pred, pred)
bridge.terminate()
return pred
def test_boosting_tree_local(self):
data = load_iris()
X = data.data
mask = np.random.choice(a=[False, True], size=X.shape, p=[0.5, 0.5])
X[mask] = float('nan')
y = np.minimum(data.target, 1)
booster = BoostingTreeEnsamble(None,
max_iters=5,
max_depth=3,
num_parallel=2)
booster.fit(X, y)
pred = booster.batch_predict(X)
self.assertGreater(sum((pred > 0.5) == y) / len(y), 0.94)
def test_boosting_tree(self):
proto = tree_pb2.BoostingTreeEnsambleProto(
params=tree_pb2.BoostingParamsProto(
num_rounds=2,
max_depth=3,
lam=1.0,
sketch_eps=0.2))
booster = BoostingTreeEnsamble(proto)
data = load_iris()
labels = data.target
labels = np.minimum(labels, 1)
features = {str(i): data.data[:, i] for i in range(data.data.shape[1])}
booster.fit(None, [i for i in range(len(data.target))], features, labels)
def run(args):
if args.verbosity == 0:
logging.basicConfig(level=logging.WARNING)
elif args.verbosity == 1:
logging.basicConfig(level=logging.INFO)
else:
logging.basicConfig(level=logging.DEBUG)
assert args.role in ['leader', 'follower', 'local'], \
"role must be leader, follower, or local"
assert args.mode in ['train', 'test', 'eval'], \
"mode must be train, test, or eval"
#follower或leader
if args.role != 'local':
bridge = Bridge(args.role, int(args.local_addr.split(':')[1]),
args.peer_addr, args.application_id, 0,
streaming_mode=args.use_streaming)
else:
bridge = None
try:
#boost
booster = BoostingTreeEnsamble(
bridge,
learning_rate=args.learning_rate,
max_iters=args.max_iters,
max_depth=args.max_depth,
l2_regularization=args.l2_regularization,
max_bins=args.max_bins,
num_parallel=args.num_parallel,
loss_type=args.loss_type,
send_scores_to_follower=args.send_scores_to_follower,
send_metrics_to_follower=args.send_metrics_to_follower)
#加载已存储的模型
if args.load_model_path:
booster.load_saved_model(args.load_model_path)
#训练不需要bridge，为什么呢
if args.mode == 'train':
train(args, booster)
#测试，评估模型需要bridge
else: # args.mode == 'test, eval'
test(args, bridge, booster)
#把模型存起来
if args.export_path:
booster.save_model(args.export_path)
except Exception as e:
logging.fatal(
'Exception raised during training: %s',
traceback.format_exc())
raise e
finally:
#结束bridge
if bridge:
bridge.terminate()