5 Best Ways to Define Feature Columns in TensorFlow with Python

Rate this post

πŸ’‘ Problem Formulation: How do you transform raw data into a set of features that TensorFlow can work with? Defining feature columns in TensorFlow is crucial when you’re preparing data for a machine learning model. Let’s say you have customer data and you want to predict churn; you’ll need to transform customer attributes into feature columns that a TensorFlow model can consume for training.

Method 1: Numeric Column

For numerical data, TensorFlow provides the numeric_column method to encapsulate the data as a numerical feature. This kind of feature column is used to represent real numbers or integers, which are directly used by the model for computation during the training process.

Here’s an example:

import tensorflow as tf

# Define a numeric feature column.
age = tf.feature_column.numeric_column('age')

# You would then add this feature column to a list of features used by a model.

Assuming your input data has an ‘age’ attribute, the model can directly use this numerical value for computation.

This example demonstrates how to create a numeric feature column for an ‘age’ attribute that could represent a customer’s age in a dataset. The created feature column can then be used by a TensorFlow estimator model for training purposes.

Method 2: Bucketized Column

To handle numerical data that should be split into different categories based on value ranges, the bucketized_column function is useful. It transforms continuous data into categorical data by segmenting it into bins.

Here’s an example:

import tensorflow as tf

# First, create a numeric column.
age = tf.feature_column.numeric_column('age')

# Then bucketize the age into different bins.
age_buckets = tf.feature_column.bucketized_column(age, boundaries=[18, 25, 30, 35, 40, 45, 50, 55, 60])

# The bucketized age column could then be used with a TensorFlow model.

Individuals will fall into buckets based on their age, with boundaries separating the different age groups.

The code snippet creates a bucketized column for the ‘age’ feature. This converts the numerical age into categories depending on the specified age boundaries, which is pivotal for models that work better with categorical data.

Method 3: Categorical Identity Column

The categorical_column_with_identity method is used when your features are categorical and represented as integers in a one-hot encoding fashion, which assigns a binary value to each category.

Here’s an example:

import tensorflow as tf

# Define a categorical identity feature column.
category = tf.feature_column.categorical_column_with_identity(key='category', num_buckets=5)

# This column could be used with TensorFlow estimators.

This will assign a unique integer in range [0, num_buckets) for each category.

In this case, a categorical identity column is established for a feature ‘category’, intended to be used when data is already in the form of integer ids with a known number of categories (num_buckets).

Method 4: Crossed Column

The crossed_column method can be used to create a feature that is a combination of two or more features, often categoricals. It’s a powerful way to capture the interaction between different feature columns.

Here’s an example:

import tensorflow as tf

# Define two categorical identity feature columns.
feature_b = tf.feature_column.categorical_column_with_identity('feature_b', num_buckets=3)
feature_c = tf.feature_column.categorical_column_with_identity('feature_c', num_buckets=4)

# Create a crossed column from the above two feature columns.
crossed_feature = tf.feature_column.crossed_column([feature_b, feature_c], hash_bucket_size=12)

# The crossed column could then be used by a model.

This creates a new column that combines features ‘feature_b’ and ‘feature_c’.

The crossed column is a feature engineering method where feature ‘feature_b’ and ‘feature_c’ are combined in a way that the interactions between these features are used by TensorFlow models, increasing the model’s capability of understanding complex relationships.

Bonus One-Liner Method 5: Embedding Column

Categorical data can be transformed into dense embeddings, which is especially useful for high-dimensional sparse data. The embedding_column takes sparse data and provides a dense, lower-dimensional representation.

Here’s an example:

import tensorflow as tf

# Start with a categorical column. 
category = tf.feature_column.categorical_column_with_vocabulary_list(
    'category', ['red', 'blue', 'green'])

# Create an embedding feature column.
embedding_feature = tf.feature_column.embedding_column(category, dimension=3)

# Use this embedding feature column in your TensorFlow model.

The result is a dense representation of the ‘category’ feature, embedded into three dimensions.

An embedding column maps a high-dimensional categorical feature into a lower-dimensional space. It’s particularly valuable for handling categorical data with many categories, translating them into input features that can be processed efficiently by neural networks.


  • Method 1: Numeric Column. Best for raw numerical data. Offers direct usage in models without the need for transformation. May not capture complex patterns well on its own.
  • Method 2: Bucketized Column. Transforms numbers into categories based on bins. Great for turning age or similarly distributed numerical features into useful categorical inputs, but might leave out nuanced information between bins.
  • Method 3: Categorical Identity Column. Maps integers to categories, perfect for one-hot encoded features. Simple and straight to the point, yet lacks the expressiveness of more nuanced encoding techniques.
  • Method 4: Crossed Column. Combines features into a single feature. Enables the model to understand complex feature interactions, but can lead to an exponential increase in dimensions needing careful tuning.
  • Bonus Method 5: Embedding Column. Converts categorical data into dense embeddings, ideal for large categorical feature sets, though it requires more computational resources and may be overkill for small datasets.