Hi Jamie,
I am copying the answer from the github 
.
As you corrected noted, the API does not allow to obtain the individual tree predictions directly. Please feel free to create a feature request :). If we see traction, we will prioritize it.
In the mean time, there is two alternative solutions:
- Training multiple Random Forest models, each with one tree (while making sure to change the random seed).
- Training a single Random Forest model and dividing it per trees using the model inspector and model builder.
Using the model builder to generate the individual trees might be easier than running the inference manually in python.
While faster than solution 1., the solution 2. can still be slow on large models and datasets as the model deserialization+re-serialization in python is relatively slow. It would look like this:
# Train a Random Forest with 10 trees
model = tfdf.keras.RandomForestModel(num_trees=10)
model.fit(train_ds)
# Extract each of the 10 trees into a separate model.
inspector = model.make_inspector()
# TODO: Run in parallel.
models = []
for tree_idx, tree in enumerate(inspector.extract_all_trees()):
print(f"Extract and export tree #{tree_idx}")
# Create a RF model with a single tree.
path = os.path.join(f"/tmp/model/{tree_idx}")
builder = tfdf.builder.RandomForestBuilder(
path=path,
objective=inspector.objective(),
import_dataspec=inspector.dataspec)
builder.add_tree(tree)
builder.close()
models.append(tf.keras.models.load_model(path))
# Compute the predictions of all the trees together.
class CombinedModel (tf.keras.Model):
def call(self, inputs):
# We assume that we have a binary classication model that returns a single
# probability. In case of multi-class classification, use tf.stack instead.
return tf.concat([ submodel(inputs) for submodel in models], axis=1)
print("Prediction of all the trees")
combined_model = CombinedModel()
all_trees_predictions = combined_model.predict(test_with_cast_ds)
See this colab for a full example.
Ps: Make sure to correctly use the all_trees_predictions to compute the prediction confidence interval. For example using Wager et al. method.
Cheers,
M.
Edit: Add the import_dataspec constructor argument in the model builder. This will help with some of the situation with categorical features. See this page for some explanations.
