TFDF customise training logs PR-AUC


I am current testing out tensorflow decision forests. I’m struggling to evaluate how the performance changes vs. number of trees for a non-default metric (in this case PR-AUC).

Below is some code with my attempts.

from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
import pandas as pd
import tensorflow as tf
import tensorflow_decision_forests as tfdf

train = load_diabetes()
X = pd.DataFrame(train['data'])
X['target'] = (pd.Series(train['target']) > 100).astype(int)
X_train, X_test = train_test_split(X)
train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(X_train, label="target")   
test_ds = tfdf.keras.pd_dataframe_to_tf_dataset(X_test, label="target")   
pr_auc = tf.keras.metrics.AUC( curve='PR',)
tfdf_clf = tfdf.keras.GradientBoostedTreesModel()
tfdf_clf.compile(metrics=[pr_auc]), validation_data=test_ds,)

Now I get very useful training logs using

#[TrainLog(num_trees=1, evaluation=Evaluation(num_examples=None, accuracy=0.9005518555641174, loss=0.6005926132202148, rmse=None, ndcg=None, aucs=None)),
#TrainLog(num_trees=2, evaluation=Evaluation(num_examples=None, accuracy=0.9005518555641174, loss=0.5672071576118469, rmse=None, ndcg=None, aucs=None)),

But it doesn’t contain any info on PR-AUC vs. iterations

If I evaluate the model, it only persists PR-AUC at the end of training.


1180/1180 [==============================] - 10s 8ms/step - loss: 0.0000e+00 - auc: 0.6832

I was wondering can anyone here help explain How I can find PR-AUC (on test data) vs. number of trees when using TFDF?

1 Like

Hello Olivier,

TF-DF has two independent means of model evaluation: 1) The internal metrics / loss used to train the model configured with the task or loss hyper-parameters, and 2) the Keras metrics supplied with compile. There is currently no connection between the two sets of metrics.

tfdf_clf.make_inspector().training_logs() returns the former while tfdf_clf.evaluate() returns the latter.

As you noted, only the internal metrics will be given for each number of trees in the model. The Keras metrics will only be given for the entire model.

Currently, the only solution to your problem is to manually remove trees from the model using the model builder and to evaluate it (more precisely, to create a new model for each number of trees you want to evaluate).

I understand that this is not a user-friendly solution. If you want, you can create a feature request on the project’s github page to help us prioritize it.



Thanks Mathieu, this is very helpful. My next question would be how do I remove x trees from my tfdf_clf model above? I have just looked online for how to do this and got a bit stuck.

Hello Olivier,

This would look as follow:

# Train a model
model = tfdf.keras.GradientBoostedTreesModel()

# Extract trees from the model.
inspector = model.make_inspector()
all_trees = inspector.extract_all_trees()

# Create a new model with only the first 5 trees.
num_extracted_trees = 5
pruned_model_path = "/tmp/pruned_model"
assert num_extracted_trees < inspector.num_trees()

builder = tfdf.builder.GradientBoostedTreeBuilder(
  path = pruned_model_path, ...)
for tree in all_trees[:num_extracted_trees]:

# Load and evaluate the new model
pruned_model = tf.keras.models.load_model(pruned_model_path)
pruned_evaluation = pruned_model.evaluate(test_ds, return_dict=True)

This notebook is a working example. It evaluates the PR-AUC vs the number of trees.


1 Like

Thank you very much for this detailed response. It is very helpful!