Watch talks from the 2019 TensorFlow Dev Summit Watch now

Build a linear model with Estimators

View on TensorFlow.org Run in Google Colab View source on GitHub

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 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.

Setup

Import TensorFlow, feature column support, and supporting modules:

import tensorflow as tf
import tensorflow.feature_column as fc 

import os
import sys

import matplotlib.pyplot as plt
from IPython.display import clear_output

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

tf.enable_eager_execution()

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:

! pip install -q requests
! git clone --depth 1 https://github.com/tensorflow/models
You are using pip version 18.1, however version 19.0.1 is available.
You should consider upgrading via the 'pip install --upgrade pip' command.
Cloning into 'models'...
remote: Enumerating objects: 2927, done.
remote: Counting objects: 100% (2927/2927), done.
remote: Compressing objects: 100% (2459/2459), done.
remote: Total 2927 (delta 520), reused 2020 (delta 393), pack-reused 0
Receiving objects: 100% (2927/2927), 369.02 MiB | 33.89 MiB/s, done.
Resolving deltas: 100% (520/520), done.
Checking connectivity... done.

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

models_path = os.path.join(os.getcwd(), 'models')

sys.path.append(models_path)

Download the dataset:

from official.wide_deep import census_dataset
from official.wide_deep import census_main

census_dataset.download("/tmp/census_data/")

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.

#export PYTHONPATH=${PYTHONPATH}:"$(pwd)/models"
#running from python you need to set the `os.environ` or the subprocess will not see the directory.

if "PYTHONPATH" in os.environ:
  os.environ['PYTHONPATH'] += os.pathsep +  models_path
else:
  os.environ['PYTHONPATH'] = models_path

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

!python -m official.wide_deep.census_main --help
2019-01-31 23:51:58.981680: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
Train DNN on census income dataset.
flags:

/docker/output/models/official/wide_deep/census_main.py:
  -bs,--batch_size:
    Batch size for training and evaluation. When using multiple gpus, this is
    the
    global batch size for all devices. For example, if the batch size is 32 and
    there are 4 GPUs, each GPU will get 8 examples on each step.
    (default: '40')
    (an integer)
  --[no]clean:
    If set, model_dir will be removed if it exists.
    (default: 'false')
  -dd,--data_dir:
    The location of the input data.
    (default: '/tmp/census_data')
  --[no]download_if_missing:
    Download data to data_dir if it is not already present.
    (default: 'true')
  -ebe,--epochs_between_evals:
    The number of training epochs to run between evaluations.
    (default: '2')
    (an integer)
  -ed,--export_dir:
    If set, a SavedModel serialization of the model will be exported to this
    directory at the end of training. See the README for more details and
    relevant
    links.
  -hk,--hooks:
    A list of (case insensitive) strings to specify the names of training hooks.
      Hook:
        loggingtensorhook
        examplespersecondhook
        loggingmetrichook
        profilerhook
      Example: `--hooks ProfilerHook,ExamplesPerSecondHook`
    See official.utils.logs.hooks_helper for details.
    (default: 'LoggingTensorHook')
    (a comma separated list)
  -md,--model_dir:
    The location of the model checkpoint files.
    (default: '/tmp/census_model')
  -mt,--model_type: <wide|deep|wide_deep>: Select model topology.
    (default: 'wide_deep')
  -te,--train_epochs:
    The number of epochs used to train.
    (default: '40')
    (an integer)

Try --helpfull to get a list of all flags.

Now run the model:

!python -m official.wide_deep.census_main --model_type=wide --train_epochs=2
2019-01-31 23:52:01.320205: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
I0131 23:52:01.344556 140274181461760 tf_logging.py:115] Using config: {'_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f93c883eb38>, '_service': None, '_num_ps_replicas': 0, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_train_distribute': None, '_tf_random_seed': None, '_eval_distribute': None, '_num_worker_replicas': 1, '_is_chief': True, '_log_step_count_steps': 100, '_task_type': 'worker', '_experimental_distribute': None, '_global_id_in_cluster': 0, '_protocol': None, '_save_checkpoints_secs': 600, '_keep_checkpoint_max': 5, '_model_dir': '/tmp/census_model', '_master': '', '_task_id': 0, '_session_config': device_count {
  key: "GPU"
}
, '_evaluation_master': '', '_device_fn': None, '_keep_checkpoint_every_n_hours': 10000}
W0131 23:52:01.345459 140274181461760 tf_logging.py:120] 'cpuinfo' not imported. CPU info will not be logged.
W0131 23:52:01.346601 140274181461760 tf_logging.py:120] 'psutil' not imported. Memory info will not be logged.
I0131 23:52:01.480358 140274181461760 tf_logging.py:115] Benchmark run: {'run_parameters': [{'long_value': 40, 'name': 'batch_size'}, {'string_value': 'wide', 'name': 'model_type'}, {'long_value': 2, 'name': 'train_epochs'}], 'tensorflow_version': {'git_hash': "b'unknown'", 'version': '1.12.0'}, 'tensorflow_environment_variables': [], 'dataset': {'name': 'Census Income'}, 'machine_config': {'gpu_info': {'count': 0}}, 'model_name': 'wide_deep', 'test_id': None, 'run_date': '2019-01-31T23:52:01.345103Z'}
I0131 23:52:01.503700 140274181461760 tf_logging.py:115] Parsing /tmp/census_data/adult.data
I0131 23:52:01.538895 140274181461760 tf_logging.py:115] Calling model_fn.
I0131 23:52:02.491058 140274181461760 tf_logging.py:115] Done calling model_fn.
I0131 23:52:02.491329 140274181461760 tf_logging.py:115] Create CheckpointSaverHook.
I0131 23:52:02.865803 140274181461760 tf_logging.py:115] Graph was finalized.
I0131 23:52:02.943673 140274181461760 tf_logging.py:115] Running local_init_op.
I0131 23:52:02.962086 140274181461760 tf_logging.py:115] Done running local_init_op.
I0131 23:52:03.667733 140274181461760 tf_logging.py:115] Saving checkpoints for 0 into /tmp/census_model/model.ckpt.
I0131 23:52:04.160730 140274181461760 tf_logging.py:115] average_loss = 0.6931472, loss = 27.725887
I0131 23:52:04.161215 140274181461760 tf_logging.py:115] loss = 27.725887, step = 1
I0131 23:52:04.715913 140274181461760 tf_logging.py:115] global_step/sec: 180.021
I0131 23:52:04.716622 140274181461760 tf_logging.py:115] average_loss = 0.409842, loss = 16.39368 (0.556 sec)
I0131 23:52:04.716863 140274181461760 tf_logging.py:115] loss = 16.39368, step = 101 (0.556 sec)
I0131 23:52:05.010882 140274181461760 tf_logging.py:115] global_step/sec: 339.058
I0131 23:52:05.011763 140274181461760 tf_logging.py:115] average_loss = 0.3900322, loss = 15.601288 (0.295 sec)
I0131 23:52:05.012072 140274181461760 tf_logging.py:115] loss = 15.601288, step = 201 (0.295 sec)
I0131 23:52:05.319190 140274181461760 tf_logging.py:115] global_step/sec: 324.35
I0131 23:52:05.320092 140274181461760 tf_logging.py:115] average_loss = 0.4894618, loss = 19.578472 (0.308 sec)
I0131 23:52:05.320424 140274181461760 tf_logging.py:115] loss = 19.578472, step = 301 (0.308 sec)
I0131 23:52:05.625391 140274181461760 tf_logging.py:115] global_step/sec: 326.571
I0131 23:52:05.626230 140274181461760 tf_logging.py:115] average_loss = 0.2718195, loss = 10.87278 (0.306 sec)
I0131 23:52:05.626539 140274181461760 tf_logging.py:115] loss = 10.87278, step = 401 (0.306 sec)
I0131 23:52:05.940114 140274181461760 tf_logging.py:115] global_step/sec: 317.772
I0131 23:52:05.941021 140274181461760 tf_logging.py:115] average_loss = 0.26650026, loss = 10.66001 (0.315 sec)
I0131 23:52:05.941363 140274181461760 tf_logging.py:115] loss = 10.66001, step = 501 (0.315 sec)
I0131 23:52:06.255339 140274181461760 tf_logging.py:115] global_step/sec: 317.227
I0131 23:52:06.256256 140274181461760 tf_logging.py:115] average_loss = 0.480035, loss = 19.2014 (0.315 sec)
I0131 23:52:06.256651 140274181461760 tf_logging.py:115] loss = 19.2014, step = 601 (0.315 sec)
I0131 23:52:06.564699 140274181461760 tf_logging.py:115] global_step/sec: 323.233
I0131 23:52:06.565562 140274181461760 tf_logging.py:115] average_loss = 0.29760113, loss = 11.904045 (0.309 sec)
I0131 23:52:06.565936 140274181461760 tf_logging.py:115] loss = 11.904045, step = 701 (0.309 sec)
I0131 23:52:06.875358 140274181461760 tf_logging.py:115] global_step/sec: 321.92
I0131 23:52:06.876131 140274181461760 tf_logging.py:115] average_loss = 0.50532234, loss = 20.212893 (0.311 sec)
I0131 23:52:06.876380 140274181461760 tf_logging.py:115] loss = 20.212893, step = 801 (0.310 sec)
I0131 23:52:07.218960 140274181461760 tf_logging.py:115] global_step/sec: 291.012
I0131 23:52:07.219774 140274181461760 tf_logging.py:115] average_loss = 0.41217932, loss = 16.487173 (0.344 sec)
I0131 23:52:07.220084 140274181461760 tf_logging.py:115] loss = 16.487173, step = 901 (0.344 sec)
I0131 23:52:07.519645 140274181461760 tf_logging.py:115] global_step/sec: 332.544
I0131 23:52:07.520327 140274181461760 tf_logging.py:115] average_loss = 0.31510335, loss = 12.604134 (0.301 sec)
I0131 23:52:07.520567 140274181461760 tf_logging.py:115] loss = 12.604134, step = 1001 (0.300 sec)
I0131 23:52:07.794004 140274181461760 tf_logging.py:115] global_step/sec: 364.517
I0131 23:52:07.794810 140274181461760 tf_logging.py:115] average_loss = 0.25020695, loss = 10.008278 (0.274 sec)
I0131 23:52:07.795114 140274181461760 tf_logging.py:115] loss = 10.008278, step = 1101 (0.275 sec)
I0131 23:52:08.094756 140274181461760 tf_logging.py:115] global_step/sec: 332.533
I0131 23:52:08.095464 140274181461760 tf_logging.py:115] average_loss = 0.2804944, loss = 11.219776 (0.301 sec)
I0131 23:52:08.095691 140274181461760 tf_logging.py:115] loss = 11.219776, step = 1201 (0.301 sec)
I0131 23:52:08.376868 140274181461760 tf_logging.py:115] global_step/sec: 354.438
I0131 23:52:08.377696 140274181461760 tf_logging.py:115] average_loss = 0.35100263, loss = 14.040106 (0.282 sec)
I0131 23:52:08.377997 140274181461760 tf_logging.py:115] loss = 14.040106, step = 1301 (0.282 sec)
I0131 23:52:08.688684 140274181461760 tf_logging.py:115] global_step/sec: 320.683
I0131 23:52:08.689425 140274181461760 tf_logging.py:115] average_loss = 0.29088816, loss = 11.635527 (0.312 sec)
I0131 23:52:08.689726 140274181461760 tf_logging.py:115] loss = 11.635527, step = 1401 (0.312 sec)
I0131 23:52:08.987321 140274181461760 tf_logging.py:115] global_step/sec: 334.868
I0131 23:52:08.988143 140274181461760 tf_logging.py:115] average_loss = 0.3747055, loss = 14.98822 (0.299 sec)
I0131 23:52:08.988460 140274181461760 tf_logging.py:115] loss = 14.98822, step = 1501 (0.299 sec)
I0131 23:52:09.285886 140274181461760 tf_logging.py:115] global_step/sec: 334.898
I0131 23:52:09.286566 140274181461760 tf_logging.py:115] average_loss = 0.4276486, loss = 17.105944 (0.298 sec)
I0131 23:52:09.286838 140274181461760 tf_logging.py:115] loss = 17.105944, step = 1601 (0.298 sec)
I0131 23:52:09.364256 140274181461760 tf_logging.py:115] Saving checkpoints for 1629 into /tmp/census_model/model.ckpt.
I0131 23:52:09.505049 140274181461760 tf_logging.py:115] Loss for final step: 0.30121216.
I0131 23:52:09.514909 140274181461760 tf_logging.py:115] Parsing /tmp/census_data/adult.test
I0131 23:52:09.543777 140274181461760 tf_logging.py:115] Calling model_fn.
W0131 23:52:10.625160 140274181461760 tf_logging.py:125] Trapezoidal rule is known to produce incorrect PR-AUCs; please switch to "careful_interpolation" instead.
W0131 23:52:10.643238 140274181461760 tf_logging.py:125] Trapezoidal rule is known to produce incorrect PR-AUCs; please switch to "careful_interpolation" instead.
I0131 23:52:10.661432 140274181461760 tf_logging.py:115] Done calling model_fn.
I0131 23:52:10.680510 140274181461760 tf_logging.py:115] Starting evaluation at 2019-01-31-23:52:10
I0131 23:52:10.799108 140274181461760 tf_logging.py:115] Graph was finalized.
I0131 23:52:10.800659 140274181461760 tf_logging.py:115] Restoring parameters from /tmp/census_model/model.ckpt-1629
I0131 23:52:10.854243 140274181461760 tf_logging.py:115] Running local_init_op.
I0131 23:52:10.889532 140274181461760 tf_logging.py:115] Done running local_init_op.
I0131 23:52:12.525223 140274181461760 tf_logging.py:115] Finished evaluation at 2019-01-31-23:52:12
I0131 23:52:12.525509 140274181461760 tf_logging.py:115] Saving dict for global step 1629: accuracy = 0.8363123, accuracy_baseline = 0.76377374, auc = 0.8840882, auc_precision_recall = 0.6959759, average_loss = 0.35078898, global_step = 1629, label/mean = 0.23622628, loss = 13.998028, precision = 0.69572425, prediction/mean = 0.23370737, recall = 0.5457618
I0131 23:52:12.744362 140274181461760 tf_logging.py:115] Saving 'checkpoint_path' summary for global step 1629: /tmp/census_model/model.ckpt-1629
I0131 23:52:12.744953 140274181461760 tf_logging.py:115] Results at epoch 2 / 2
I0131 23:52:12.745047 140274181461760 tf_logging.py:115] ------------------------------------------------------------
I0131 23:52:12.745147 140274181461760 tf_logging.py:115] accuracy: 0.8363123
I0131 23:52:12.745221 140274181461760 tf_logging.py:115] accuracy_baseline: 0.76377374
I0131 23:52:12.745313 140274181461760 tf_logging.py:115] auc: 0.8840882
I0131 23:52:12.745393 140274181461760 tf_logging.py:115] auc_precision_recall: 0.6959759
I0131 23:52:12.745459 140274181461760 tf_logging.py:115] average_loss: 0.35078898
I0131 23:52:12.745575 140274181461760 tf_logging.py:115] global_step: 1629
I0131 23:52:12.745649 140274181461760 tf_logging.py:115] label/mean: 0.23622628
I0131 23:52:12.745724 140274181461760 tf_logging.py:115] loss: 13.998028
I0131 23:52:12.745814 140274181461760 tf_logging.py:115] precision: 0.69572425
I0131 23:52:12.745880 140274181461760 tf_logging.py:115] prediction/mean: 0.23370737
I0131 23:52:12.745964 140274181461760 tf_logging.py:115] recall: 0.5457618
I0131 23:52:12.746120 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.8363122940063477, 'timestamp': '2019-01-31T23:52:12.746081Z', 'name': 'accuracy', 'global_step': 1629}
I0131 23:52:12.746267 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.7637737393379211, 'timestamp': '2019-01-31T23:52:12.746244Z', 'name': 'accuracy_baseline', 'global_step': 1629}
I0131 23:52:12.746391 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.8840882182121277, 'timestamp': '2019-01-31T23:52:12.746361Z', 'name': 'auc', 'global_step': 1629}
I0131 23:52:12.746510 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.6959758996963501, 'timestamp': '2019-01-31T23:52:12.746490Z', 'name': 'auc_precision_recall', 'global_step': 1629}
I0131 23:52:12.746638 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.35078898072242737, 'timestamp': '2019-01-31T23:52:12.746609Z', 'name': 'average_loss', 'global_step': 1629}
I0131 23:52:12.746749 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.23622627556324005, 'timestamp': '2019-01-31T23:52:12.746729Z', 'name': 'label/mean', 'global_step': 1629}
I0131 23:52:12.746875 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 13.998027801513672, 'timestamp': '2019-01-31T23:52:12.746852Z', 'name': 'loss', 'global_step': 1629}
I0131 23:52:12.746981 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.6957242488861084, 'timestamp': '2019-01-31T23:52:12.746962Z', 'name': 'precision', 'global_step': 1629}
I0131 23:52:12.747106 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.23370736837387085, 'timestamp': '2019-01-31T23:52:12.747085Z', 'name': 'prediction/mean', 'global_step': 1629}
I0131 23:52:12.747211 140274181461760 tf_logging.py:115] Benchmark metric: {'extras': [], 'unit': None, 'value': 0.5457618236541748, 'timestamp': '2019-01-31T23:52:12.747192Z', 'name': 'recall', 'global_step': 1629}

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:

!ls  /tmp/census_data/
adult.data  adult.test
train_file = "/tmp/census_data/adult.data"
test_file = "/tmp/census_data/adult.test"

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

import pandas

train_df = pandas.read_csv(train_file, header = None, names = census_dataset._CSV_COLUMNS)
test_df = pandas.read_csv(test_file, header = None, names = census_dataset._CSV_COLUMNS)

train_df.head()
age workclass fnlwgt education education_num marital_status occupation relationship race gender capital_gain capital_loss hours_per_week native_country income_bracket
0 39 State-gov 77516 Bachelors 13 Never-married Adm-clerical Not-in-family White Male 2174 0 40 United-States <=50K
1 50 Self-emp-not-inc 83311 Bachelors 13 Married-civ-spouse Exec-managerial Husband White Male 0 0 13 United-States <=50K
2 38 Private 215646 HS-grad 9 Divorced Handlers-cleaners Not-in-family White Male 0 0 40 United-States <=50K
3 53 Private 234721 11th 7 Married-civ-spouse Handlers-cleaners Husband Black Male 0 0 40 United-States <=50K
4 28 Private 338409 Bachelors 13 Married-civ-spouse Prof-specialty Wife Black Female 0 0 40 Cuba <=50K

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.

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

def easy_input_function(df, label_key, num_epochs, shuffle, batch_size):
  label = df[label_key]
  ds = tf.data.Dataset.from_tensor_slices((dict(df),label))

  if shuffle:
    ds = ds.shuffle(10000)

  ds = ds.batch(batch_size).repeat(num_epochs)

  return ds

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

ds = easy_input_function(train_df, label_key='income_bracket', num_epochs=5, shuffle=True, batch_size=10)

for feature_batch, label_batch in ds.take(1):
  print('Some feature keys:', list(feature_batch.keys())[:5])
  print()
  print('A batch of Ages  :', feature_batch['age'])
  print()
  print('A batch of Labels:', label_batch )
Some feature keys: ['education', 'occupation', 'capital_loss', 'gender', 'age']

A batch of Ages  : tf.Tensor([47 24 37 32 78 20 53 42 60 30], shape=(10,), dtype=int32)

A batch of Labels: tf.Tensor(
[b'<=50K' b'<=50K' b'<=50K' b'<=50K' b'>50K' b'<=50K' b'<=50K' b'>50K'
 b'<=50K' b'>50K'], shape=(10,), dtype=string)

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:

import inspect
print(inspect.getsource(census_dataset.input_fn))
def input_fn(data_file, num_epochs, shuffle, batch_size):
  """Generate an input function for the Estimator."""
  assert tf.gfile.Exists(data_file), (
      '%s not found. Please make sure you have run census_dataset.py and '
      'set the --data_dir argument to the correct path.' % data_file)

  def parse_csv(value):
    tf.logging.info('Parsing {}'.format(data_file))
    columns = tf.decode_csv(value, record_defaults=_CSV_COLUMN_DEFAULTS)
    features = dict(zip(_CSV_COLUMNS, columns))
    labels = features.pop('income_bracket')
    classes = tf.equal(labels, '>50K')  # binary classification
    return features, classes

  # Extract lines from input files using the Dataset API.
  dataset = tf.data.TextLineDataset(data_file)

  if shuffle:
    dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'])

  dataset = dataset.map(parse_csv, num_parallel_calls=5)

  # We call repeat after shuffling, rather than before, to prevent separate
  # epochs from blending together.
  dataset = dataset.repeat(num_epochs)
  dataset = dataset.batch(batch_size)
  return dataset

This input_fn returns equivalent output:

ds = census_dataset.input_fn(train_file, num_epochs=5, shuffle=True, batch_size=10)

for feature_batch, label_batch in ds.take(1):
  print('Feature keys:', list(feature_batch.keys())[:5])
  print()
  print('Age batch   :', feature_batch['age'])
  print()
  print('Label batch :', label_batch )
INFO:tensorflow:Parsing /tmp/census_data/adult.data

WARNING: Logging before flag parsing goes to stderr.
I0131 23:52:14.567273 140073592747776 tf_logging.py:115] Parsing /tmp/census_data/adult.data

Feature keys: ['education', 'occupation', 'capital_loss', 'gender', 'age']

Age batch   : tf.Tensor([64 90 20 61 22 25 19 36 28 21], shape=(10,), dtype=int32)

Label batch : tf.Tensor([False False False False False False False False False False], shape=(10,), dtype=bool)

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

import functools

train_inpf = functools.partial(census_dataset.input_fn, train_file, num_epochs=2, shuffle=True, batch_size=64)
test_inpf = functools.partial(census_dataset.input_fn, test_file, num_epochs=1, shuffle=False, batch_size=64)

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:

age = fc.numeric_column('age')

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

fc.input_layer(feature_batch, [age]).numpy()
array([[64.],
       [90.],
       [20.],
       [61.],
       [22.],
       [25.],
       [19.],
       [36.],
       [28.],
       [21.]], dtype=float32)

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

classifier = tf.estimator.LinearClassifier(feature_columns=[age])
classifier.train(train_inpf)
result = classifier.evaluate(test_inpf)

clear_output()  # used for display in notebook
print(result)
{'precision': 0.21177803, 'loss': 33.873966, 'average_loss': 0.5305486, 'global_step': 1018, 'auc': 0.6783597, 'accuracy': 0.73251027, 'label/mean': 0.23622628, 'prediction/mean': 0.2917571, 'auc_precision_recall': 0.31139234, 'accuracy_baseline': 0.76377374, 'recall': 0.048621945}

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

education_num = tf.feature_column.numeric_column('education_num')
capital_gain = tf.feature_column.numeric_column('capital_gain')
capital_loss = tf.feature_column.numeric_column('capital_loss')
hours_per_week = tf.feature_column.numeric_column('hours_per_week')

my_numeric_columns = [age,education_num, capital_gain, capital_loss, hours_per_week]

fc.input_layer(feature_batch, my_numeric_columns).numpy()
array([[  64.,    0.,    0.,   13.,   30.],
       [  90., 6767.,    0.,    9.,   40.],
       [  20.,    0.,    0.,   10.,   40.],
       [  61.,    0.,    0.,    9.,   40.],
       [  22.,    0.,    0.,    9.,   40.],
       [  25.,    0.,    0.,    9.,    8.],
       [  19.,    0.,    0.,   10.,   16.],
       [  36.,    0.,    0.,    9.,   35.],
       [  28.,    0.,    0.,   13.,   25.],
       [  21.,    0.,    0.,    9.,   35.]], dtype=float32)

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

classifier = tf.estimator.LinearClassifier(feature_columns=my_numeric_columns)
classifier.train(train_inpf)

result = classifier.evaluate(test_inpf)

clear_output()

for key,value in sorted(result.items()):
  print('%s: %s' % (key, value))
accuracy: 0.781893
accuracy_baseline: 0.76377374
auc: 0.7717221
auc_precision_recall: 0.5343629
average_loss: 0.9750157
global_step: 1018
label/mean: 0.23622628
loss: 62.251884
precision: 0.5710843
prediction/mean: 0.29361963
recall: 0.30811232

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.

relationship = fc.categorical_column_with_vocabulary_list(
    'relationship',
    ['Husband', 'Not-in-family', 'Wife', 'Own-child', 'Unmarried', 'Other-relative'])

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

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

fc.input_layer(feature_batch, [age, fc.indicator_column(relationship)])
WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/sparse_ops.py:1165: sparse_to_dense (from tensorflow.python.ops.sparse_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Create a <a href="../../api_docs/python/tf/sparse/SparseTensor"><code>tf.sparse.SparseTensor</code></a> and use <a href="../../api_docs/python/tf/sparse/to_dense"><code>tf.sparse.to_dense</code></a> instead.

W0131 23:52:28.295203 140073592747776 tf_logging.py:125] From /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/sparse_ops.py:1165: sparse_to_dense (from tensorflow.python.ops.sparse_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Create a <a href="../../api_docs/python/tf/sparse/SparseTensor"><code>tf.sparse.SparseTensor</code></a> and use <a href="../../api_docs/python/tf/sparse/to_dense"><code>tf.sparse.to_dense</code></a> instead.

<tf.Tensor: id=4578, shape=(10, 7), dtype=float32, numpy=
array([[64.,  0.,  1.,  0.,  0.,  0.,  0.],
       [90.,  1.,  0.,  0.,  0.,  0.,  0.],
       [20.,  0.,  0.,  0.,  1.,  0.,  0.],
       [61.,  1.,  0.,  0.,  0.,  0.,  0.],
       [22.,  0.,  0.,  0.,  1.,  0.,  0.],
       [25.,  0.,  0.,  0.,  1.,  0.,  0.],
       [19.,  0.,  0.,  0.,  1.,  0.,  0.],
       [36.,  0.,  1.,  0.,  0.,  0.,  0.],
       [28.,  0.,  0.,  0.,  0.,  0.,  1.],
       [21.,  0.,  0.,  0.,  1.,  0.,  0.]], dtype=float32)>

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

occupation = tf.feature_column.categorical_column_with_hash_bucket(
    'occupation', hash_bucket_size=1000)

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:

for item in feature_batch['occupation'].numpy():
    print(item.decode())
Prof-specialty
Other-service
Sales
Other-service
Craft-repair
Other-service
?
Transport-moving
Handlers-cleaners
Other-service

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

occupation_result = fc.input_layer(feature_batch, [fc.indicator_column(occupation)])

occupation_result.numpy().shape
(10, 1000)

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:

tf.argmax(occupation_result, axis=1).numpy()
array([979, 527, 631, 527, 466, 527,  65, 420,  10, 527])

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:

education = tf.feature_column.categorical_column_with_vocabulary_list(
    'education', [
        'Bachelors', 'HS-grad', '11th', 'Masters', '9th', 'Some-college',
        'Assoc-acdm', 'Assoc-voc', '7th-8th', 'Doctorate', 'Prof-school',
        '5th-6th', '10th', '1st-4th', 'Preschool', '12th'])

marital_status = tf.feature_column.categorical_column_with_vocabulary_list(
    'marital_status', [
        'Married-civ-spouse', 'Divorced', 'Married-spouse-absent',
        'Never-married', 'Separated', 'Married-AF-spouse', 'Widowed'])

workclass = tf.feature_column.categorical_column_with_vocabulary_list(
    'workclass', [
        'Self-emp-not-inc', 'Private', 'State-gov', 'Federal-gov',
        'Local-gov', '?', 'Self-emp-inc', 'Without-pay', 'Never-worked'])


my_categorical_columns = [relationship, occupation, education, marital_status, workclass]

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

classifier = tf.estimator.LinearClassifier(feature_columns=my_numeric_columns+my_categorical_columns)
classifier.train(train_inpf)
result = classifier.evaluate(test_inpf)

clear_output()

for key,value in sorted(result.items()):
  print('%s: %s' % (key, value))
accuracy: 0.8266077
accuracy_baseline: 0.76377374
auc: 0.8773386
auc_precision_recall: 0.66341794
average_loss: 0.78451854
global_step: 1018
label/mean: 0.23622628
loss: 50.0892
precision: 0.63746303
prediction/mean: 0.2577103
recall: 0.6167447

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:

age_buckets = tf.feature_column.bucketized_column(
    age, boundaries=[18, 25, 30, 35, 40, 45, 50, 55, 60, 65])

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 a one-hot feature:

fc.input_layer(feature_batch, [age, age_buckets]).numpy()
array([[64.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [90.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.],
       [20.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [61.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [22.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [25.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [19.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [36.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.],
       [28.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [21.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]],
      dtype=float32)

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:

education_x_occupation = tf.feature_column.crossed_column(
    ['education', 'occupation'], hash_bucket_size=1000)

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:

age_buckets_x_education_x_occupation = tf.feature_column.crossed_column(
    [age_buckets, 'education', 'occupation'], hash_bucket_size=1000)

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.

For a visual example the effect of hash-buckets with crossed columns see this notebook

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:

import tempfile

base_columns = [
    education, marital_status, relationship, workclass, occupation,
    age_buckets,
]

crossed_columns = [
    tf.feature_column.crossed_column(
        ['education', 'occupation'], hash_bucket_size=1000),
    tf.feature_column.crossed_column(
        [age_buckets, 'education', 'occupation'], hash_bucket_size=1000),
]

model = tf.estimator.LinearClassifier(
    model_dir=tempfile.mkdtemp(), 
    feature_columns=base_columns + crossed_columns,
    optimizer=tf.train.FtrlOptimizer(learning_rate=0.1))
INFO:tensorflow:Using default config.

I0131 23:52:37.807945 140073592747776 tf_logging.py:115] Using default config.

INFO:tensorflow:Using config: {'_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_save_checkpoints_secs': 600, '_is_chief': True, '_protocol': None, '_save_checkpoints_steps': None, '_global_id_in_cluster': 0, '_master': '', '_eval_distribute': None, '_task_type': 'worker', '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f65057b44a8>, '_evaluation_master': '', '_model_dir': '/tmp/tmp4rk7grt4', '_tf_random_seed': None, '_device_fn': None, '_experimental_distribute': None, '_keep_checkpoint_max': 5, '_num_worker_replicas': 1, '_service': None, '_save_summary_steps': 100, '_task_id': 0, '_num_ps_replicas': 0}

I0131 23:52:37.810586 140073592747776 tf_logging.py:115] Using config: {'_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_save_checkpoints_secs': 600, '_is_chief': True, '_protocol': None, '_save_checkpoints_steps': None, '_global_id_in_cluster': 0, '_master': '', '_eval_distribute': None, '_task_type': 'worker', '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f65057b44a8>, '_evaluation_master': '', '_model_dir': '/tmp/tmp4rk7grt4', '_tf_random_seed': None, '_device_fn': None, '_experimental_distribute': None, '_keep_checkpoint_max': 5, '_num_worker_replicas': 1, '_service': None, '_save_summary_steps': 100, '_task_id': 0, '_num_ps_replicas': 0}

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:

train_inpf = functools.partial(census_dataset.input_fn, train_file, 
                               num_epochs=40, shuffle=True, batch_size=64)

model.train(train_inpf)

clear_output()  # used for notebook display

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

results = model.evaluate(test_inpf)

clear_output()

for key,value in sorted(result.items()):
  print('%s: %0.2f' % (key, value))
accuracy: 0.83
accuracy_baseline: 0.76
auc: 0.88
auc_precision_recall: 0.66
average_loss: 0.78
global_step: 1018.00
label/mean: 0.24
loss: 50.09
precision: 0.64
prediction/mean: 0.26
recall: 0.62

The first line of the output should display something like: accuracy: 0.83, which means the accuracy is 83%. 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:

import numpy as np

predict_df = test_df[:20].copy()

pred_iter = model.predict(
    lambda:easy_input_function(predict_df, label_key='income_bracket',
                               num_epochs=1, shuffle=False, batch_size=10))

classes = np.array(['<=50K', '>50K'])
pred_class_id = []

for pred_dict in pred_iter:
  pred_class_id.append(pred_dict['class_ids'])

predict_df['predicted_class'] = classes[np.array(pred_class_id)]
predict_df['correct'] = predict_df['predicted_class'] == predict_df['income_bracket']

clear_output()

predict_df[['income_bracket','predicted_class', 'correct']]
income_bracket predicted_class correct
0 <=50K <=50K True
1 <=50K <=50K True
2 >50K <=50K False
3 >50K <=50K False
4 <=50K <=50K True
5 <=50K <=50K True
6 <=50K <=50K True
7 >50K >50K True
8 <=50K <=50K True
9 <=50K <=50K True
10 >50K <=50K False
11 <=50K >50K False
12 <=50K <=50K True
13 <=50K <=50K True
14 >50K <=50K False
15 >50K >50K True
16 <=50K <=50K True
17 <=50K <=50K True
18 <=50K <=50K True
19 >50K >50K True

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:

model_l1 = tf.estimator.LinearClassifier(
    feature_columns=base_columns + crossed_columns,
    optimizer=tf.train.FtrlOptimizer(
        learning_rate=0.1,
        l1_regularization_strength=10.0,
        l2_regularization_strength=0.0))

model_l1.train(train_inpf)

results = model_l1.evaluate(test_inpf)
clear_output()
for key in sorted(results):
  print('%s: %0.2f' % (key, results[key]))
accuracy: 0.84
accuracy_baseline: 0.76
auc: 0.88
auc_precision_recall: 0.69
average_loss: 0.35
global_step: 20351.00
label/mean: 0.24
loss: 22.47
precision: 0.69
prediction/mean: 0.24
recall: 0.56
model_l2 = tf.estimator.LinearClassifier(
    feature_columns=base_columns + crossed_columns,
    optimizer=tf.train.FtrlOptimizer(
        learning_rate=0.1,
        l1_regularization_strength=0.0,
        l2_regularization_strength=10.0))

model_l2.train(train_inpf)

results = model_l2.evaluate(test_inpf)
clear_output()
for key in sorted(results):
  print('%s: %0.2f' % (key, results[key]))
accuracy: 0.84
accuracy_baseline: 0.76
auc: 0.88
auc_precision_recall: 0.69
average_loss: 0.35
global_step: 20351.00
label/mean: 0.24
loss: 22.46
precision: 0.69
prediction/mean: 0.24
recall: 0.55

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:

def get_flat_weights(model):
  weight_names = [
      name for name in model.get_variable_names()
      if "linear_model" in name and "Ftrl" not in name]

  weight_values = [model.get_variable_value(name) for name in weight_names]

  weights_flat = np.concatenate([item.flatten() for item in weight_values], axis=0)

  return weights_flat

weights_flat = get_flat_weights(model)
weights_flat_l1 = get_flat_weights(model_l1)
weights_flat_l2 = get_flat_weights(model_l2)

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:

weight_mask = weights_flat != 0

weights_base = weights_flat[weight_mask]
weights_l1 = weights_flat_l1[weight_mask]
weights_l2 = weights_flat_l2[weight_mask]

Now plot the distributions:

plt.figure()
_ = plt.hist(weights_base, bins=np.linspace(-3,3,30))
plt.title('Base Model')
plt.ylim([0,500])

plt.figure()
_ = plt.hist(weights_l1, bins=np.linspace(-3,3,30))
plt.title('L1 - Regularization')
plt.ylim([0,500])

plt.figure()
_ = plt.hist(weights_l2, bins=np.linspace(-3,3,30))
plt.title('L2 - Regularization')
_=plt.ylim([0,500])

png

png

png

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.