gbdt_embedding.py

import sys
import numpy as np

from sklearn.datasets import load_svmlight_file
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.externals.joblib import Parallel, delayed

from sklearn.utils import array2d
from sklearn.tree._tree import DTYPE

X_train, y_train = load_svmlight_file(sys.argv[1])
X_train = X_train.toarray()

def my_func(tree, X):
    return tree.apply(array2d(X, dtype=DTYPE))

clf = GradientBoostingClassifier(n_estimators=int(sys.argv[2]), max_depth=int(sys.argv[3])).fit(X_train, y_train)
out = Parallel(n_jobs=1)(delayed(my_func)(clf.estimators_[i, 0].tree_, X_train) for i in xrange(clf.n_estimators))
out = np.transpose(np.array(out))
print out
print out.shape