Copyright 2018 The TensorFlow Authors.

Build a linear model with Estimators

Note: This is an archived TF1 notebook. These are configured to run in TF2's compatbility mode but will run in TF1 as well. To use TF1 in Colab, use the magic.

This tutorial uses the tf.estimator API in TensorFlow to solve a benchmark binary classification problem. Estimators are TensorFlow's most scalable and production-oriented model type. For more information see the Estimator guide.

Overview

Using census data which contains data about a person's age, education, marital status, and occupation (the features), we will try to predict whether or not the person earns more than 50,000 dollars a year (the target label). We will train a logistic regression model that, given an individual's information, outputs a number between 0 and 1—this can be interpreted as the probability that the individual has an annual income of over 50,000 dollars.

Key Point: As a modeler and developer, think about how this data is used and the potential benefits and harm a model's predictions can cause. A model like this could reinforce societal biases and disparities. Is each feature relevant to the problem you want to solve or will it introduce bias? For more information, read about ML fairness.

Setup

Import TensorFlow, feature column support, and supporting modules:

And let's enable eager execution to inspect this program as we run it:

Download the official implementation

We'll use the wide and deep model available in TensorFlow's model repository. Download the code, add the root directory to your Python path, and jump to the wide_deep directory:

Add the root directory of the repository to your Python path:

Download the dataset:

Command line usage

The repo includes a complete program for experimenting with this type of model.

To execute the tutorial code from the command line first add the path to tensorflow/models to your PYTHONPATH.

Use --help to see what command line options are available:

Now run the model:

Read the U.S. Census data

This example uses the U.S Census Income Dataset from 1994 and 1995. We have provided the census_dataset.py script to download the data and perform a little cleanup.

Since the task is a binary classification problem, we'll construct a label column named "label" whose value is 1 if the income is over 50K, and 0 otherwise. For reference, see the input_fn in census_main.py.

Let's look at the data to see which columns we can use to predict the target label:

pandas provides some convenient utilities for data analysis. Here's a list of columns available in the Census Income dataset:

The columns are grouped into two types: categorical and continuous columns:

• A column is called categorical if its value can only be one of the categories in a finite set. For example, the relationship status of a person (wife, husband, unmarried, etc.) or the education level (high school, college, etc.) are categorical columns.
• A column is called continuous if its value can be any numerical value in a continuous range. For example, the capital gain of a person (e.g. $14,084) is a continuous column.

Converting Data into Tensors

When building a tf.estimator model, the input data is specified by using an input function (or input_fn). This builder function returns a tf.data.Dataset of batches of (features-dict, label) pairs. It is not called until it is passed to tf.estimator.Estimator methods such as train and evaluate.

The input builder function returns the following pair:

1. features: A dict from feature names to Tensors or SparseTensors containing batches of features.
2. labels: A Tensor containing batches of labels.

The keys of the features are used to configure the model's input layer.

Note: The input function is called while constructing the TensorFlow graph, not while running the graph. It is returning a representation of the input data as a sequence of TensorFlow graph operations.

For small problems like this, it's easy to make a tf.data.Dataset by slicing the pandas.DataFrame:

Since we have eager execution enabled, it's easy to inspect the resulting dataset:

But this approach has severly-limited scalability. Larger datasets should be streamed from disk. The census_dataset.input_fn provides an example of how to do this using tf.decode_csv and tf.data.TextLineDataset:

This input_fn returns equivalent output:

Because Estimators expect an input_fn that takes no arguments, we typically wrap configurable input function into an object with the expected signature. For this notebook configure the train_inpf to iterate over the data twice:

Selecting and Engineering Features for the Model

Estimators use a system called feature columns to describe how the model should interpret each of the raw input features. An Estimator expects a vector of numeric inputs, and feature columns describe how the model should convert each feature.

Selecting and crafting the right set of feature columns is key to learning an effective model. A feature column can be either one of the raw inputs in the original features dict (a base feature column), or any new columns created using transformations defined over one or multiple base columns (a derived feature columns).

A feature column is an abstract concept of any raw or derived variable that can be used to predict the target label.

Base Feature Columns

Numeric columns

The simplest feature_column is numeric_column. This indicates that a feature is a numeric value that should be input to the model directly. For example:

The model will use the feature_column definitions to build the model input. You can inspect the resulting output using the input_layer function:

The following will train and evaluate a model using only the age feature:

Similarly, we can define a NumericColumn for each continuous feature column that we want to use in the model:

You could retrain a model on these features by changing the feature_columns argument to the constructor:

Categorical columns

To define a feature column for a categorical feature, create a CategoricalColumn using one of the tf.feature_column.categorical_column* functions.

If you know the set of all possible feature values of a column—and there are only a few of them—use categorical_column_with_vocabulary_list. Each key in the list is assigned an auto-incremented ID starting from 0. For example, for the relationship column we can assign the feature string Husband to an integer ID of 0 and "Not-in-family" to 1, etc.

This creates a sparse one-hot vector from the raw input feature.

The input_layer function we're using is designed for DNN models and expects dense inputs. To demonstrate the categorical column we must wrap it in a tf.feature_column.indicator_column to create the dense one-hot output (Linear Estimators can often skip this dense-step).

Note: the other sparse-to-dense option is tf.feature_column.embedding_column.

Run the input layer, configured with both the age and relationship columns:

If we don't know the set of possible values in advance, use the categorical_column_with_hash_bucket instead:

Here, each possible value in the feature column occupation is hashed to an integer ID as we encounter them in training. The example batch has a few different occupations:

If we run input_layer with the hashed column, we see that the output shape is (batch_size, hash_bucket_size):

It's easier to see the actual results if we take the tf.argmax over the hash_bucket_size dimension. Notice how any duplicate occupations are mapped to the same pseudo-random index:

Note: Hash collisions are unavoidable, but often have minimal impact on model quality. The effect may be noticable if the hash buckets are being used to compress the input space.

No matter how we choose to define a SparseColumn, each feature string is mapped into an integer ID by looking up a fixed mapping or by hashing. Under the hood, the LinearModel class is responsible for managing the mapping and creating tf.Variable to store the model parameters (model weights) for each feature ID. The model parameters are learned through the model training process described later.

Let's do the similar trick to define the other categorical features:

It's easy to use both sets of columns to configure a model that uses all these features:

Derived feature columns

Make Continuous Features Categorical through Bucketization

Sometimes the relationship between a continuous feature and the label is not linear. For example, age and income—a person's income may grow in the early stage of their career, then the growth may slow at some point, and finally, the income decreases after retirement. In this scenario, using the raw age as a real-valued feature column might not be a good choice because the model can only learn one of the three cases:

1. Income always increases at some rate as age grows (positive correlation),
2. Income always decreases at some rate as age grows (negative correlation), or
3. Income stays the same no matter at what age (no correlation).

If we want to learn the fine-grained correlation between income and each age group separately, we can leverage bucketization. Bucketization is a process of dividing the entire range of a continuous feature into a set of consecutive buckets, and then converting the original numerical feature into a bucket ID (as a categorical feature) depending on which bucket that value falls into. So, we can define a bucketized_column over age as:

boundaries is a list of bucket boundaries. In this case, there are 10 boundaries, resulting in 11 age group buckets (from age 17 and below, 18-24, 25-29, ..., to 65 and over).

With bucketing, the model sees each bucket as a one-hot feature:

Learn complex relationships with crossed column

Using each base feature column separately may not be enough to explain the data. For example, the correlation between education and the label (earning > 50,000 dollars) may be different for different occupations. Therefore, if we only learn a single model weight for education="Bachelors" and education="Masters", we won't capture every education-occupation combination (e.g. distinguishing between education="Bachelors" AND occupation="Exec-managerial" AND education="Bachelors" AND occupation="Craft-repair").

To learn the differences between different feature combinations, we can add crossed feature columns to the model:

We can also create a crossed_column over more than two columns. Each constituent column can be either a base feature column that is categorical (SparseColumn), a bucketized real-valued feature column, or even another CrossColumn. For example:

These crossed columns always use hash buckets to avoid the exponential explosion in the number of categories, and put the control over number of model weights in the hands of the user.

Define the logistic regression model

After processing the input data and defining all the feature columns, we can put them together and build a logistic regression model. The previous section showed several types of base and derived feature columns, including:

• CategoricalColumn
• NumericColumn
• BucketizedColumn
• CrossedColumn

All of these are subclasses of the abstract FeatureColumn class and can be added to the feature_columns field of a model:

The model automatically learns a bias term, which controls the prediction made without observing any features. The learned model files are stored in model_dir.

Train and evaluate the model

After adding all the features to the model, let's train the model. Training a model is just a single command using the tf.estimator API:

After the model is trained, evaluate the accuracy of the model by predicting the labels of the holdout data:

The first line of the output should display something like: accuracy: 0.84, which means the accuracy is 84%. You can try using more features and transformations to see if you can do better!

After the model is evaluated, we can use it to predict whether an individual has an annual income of over 50,000 dollars given an individual's information input.

Let's look in more detail how the model performed:

For a working end-to-end example, download our example code and set the model_type flag to wide.

Adding Regularization to Prevent Overfitting

Regularization is a technique used to avoid overfitting. Overfitting happens when a model performs well on the data it is trained on, but worse on test data that the model has not seen before. Overfitting can occur when a model is excessively complex, such as having too many parameters relative to the number of observed training data. Regularization allows you to control the model's complexity and make the model more generalizable to unseen data.

You can add L1 and L2 regularizations to the model with the following code:

These regularized models don't perform much better than the base model. Let's look at the model's weight distributions to better see the effect of the regularization:

The models have many zero-valued weights caused by unused hash bins (there are many more hash bins than categories in some columns). We can mask these weights when viewing the weight distributions:

Now plot the distributions:

Both types of regularization squeeze the distribution of weights towards zero. L2 regularization has a greater effect in the tails of the distribution eliminating extreme weights. L1 regularization produces more exactly-zero values, in this case it sets ~200 to zero.