decision-forests Retrieving standard deviation of predictions from ensemble

I'm working on an application where I'd like to retrieve the standard deviation of the predictions made by the trees within an ensemble (currently a tfdf.keras.RandomForestModel) to use as an estimate of the confidence of a given prediction.

It looks like I could do this by running a prediction on each individual tree with inspector.iterate_on_nodes() but is there a better way to do this via the main predict method, and if not would you consider this as an enhancement?

Jul 27 '21 08:07 jamiecollinson

Hi Jamie,

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 new 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())
  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.

Cheers, M.