TensorFlow Decision Forests

1.12.0 · active · verified Mon Apr 13

TensorFlow Decision Forests (TF-DF) is a Python library that integrates state-of-the-art decision forest algorithms (such as Random Forests and Gradient Boosted Trees) directly into TensorFlow and Keras. It enables training, serving, and interpreting these models for classification, regression, and ranking tasks. Built on the highly optimized C++ Yggdrasil Decision Forests (YDF) library, it is actively maintained with frequent releases, typically every few months.

Warnings

breaking TensorFlow Version Compatibility: Each TF-DF version is strictly tied to a specific TensorFlow version due to ABI compatibility. Using incompatible versions will lead to cryptic C++ runtime errors (e.g., 'undefined symbol').
Fix: Always check the official TF-DF documentation's compatibility table for the exact TensorFlow version required for your TF-DF version. For TF-DF 1.12.0, TensorFlow 2.19.0 is required.
deprecated Loss function `LAMBDA_MART_NDCG5` has been renamed to `LAMBDA_MART_NDCG`.
Fix: Update your code to use `LAMBDA_MART_NDCG`. The old name is still available for backward compatibility but is deprecated.
gotcha Keras 3 Incompatibility: TF-DF is not yet compatible with Keras 3.
Fix: Use `tf_keras` instead of `tf.keras`, or ensure your TensorFlow version is prior to 2.16. Many examples include `os.environ['TF_USE_LEGACY_KERAS'] = '1'` to force Keras 2 compatibility.
gotcha Windows Support Limitations: A native Windows Pip package is not available.
Fix: Windows users should utilize Windows Subsystem for Linux (WSL) and follow the Linux installation instructions.
gotcha No GPU/TPU Support: TF-DF models currently do not leverage GPUs or TPUs for training or inference.
Fix: TF-DF is designed for CPU-based training, which is often efficient for tabular data. If GPU acceleration is critical, consider alternative libraries or migrating to YDF which may offer different integration points.
gotcha Recommendation to migrate to Yggdrasil Decision Forests (YDF) for new projects.
Fix: The TF-DF team recommends YDF for new projects as it offers more features, a simplified API, and faster training times. TF-DF models are compatible with YDF models.

Install

pip install tensorflow_decision_forests Basic Installation
pip install tensorflow_decision_forests wurlitzer Installation with verbose logs (e.g., Colab)

Imports

tensorflow_decision_forests

import tensorflow_decision_forests as tfdf

RandomForestModel
```
tfdf.keras.RandomForestModel
```
All TF-DF models are Keras models and reside under `tfdf.keras`.
pd_dataframe_to_tf_dataset
```
tfdf.keras.pd_dataframe_to_tf_dataset
```
Data utility functions are also part of the `tfdf.keras` submodule.

Quickstart

This quickstart demonstrates how to train a basic RandomForestModel using TF-DF. It covers importing necessary libraries, preparing data using `pd_dataframe_to_tf_dataset`, and training the model. Ensure `pandas` is installed for this example.

import os
import tensorflow as tf
import tensorflow_decision_forests as tfdf
import pandas as pd

# Ensure Keras 2 compatibility, often needed in TF-DF environments
os.environ['TF_USE_LEGACY_KERAS'] = '1'

# Create a dummy dataset (replace with your actual data loading)
data = {
    'feature_1': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'feature_2': ['A', 'B', 'A', 'C', 'B', 'A', 'C', 'B', 'A', 'C'],
    'label': [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
}
df = pd.DataFrame(data)

# Convert the Pandas DataFrame to a TensorFlow Dataset
train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(df, label='label')

# Create and train a Random Forest model
model = tfdf.keras.RandomForestModel()
model.fit(train_ds)

# (Optional) Evaluate the model
# model.evaluate(train_ds) # Use a separate test_ds for proper evaluation

print("Model trained successfully!")
model.summary()

view raw JSON →