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.0, however version 18.1 is available.
You should consider upgrading via the 'pip install --upgrade pip' command.
Cloning into 'models'...
remote: Enumerating objects: 2978, done.
remote: Counting objects: 100% (2978/2978), done.
remote: Compressing objects: 100% (2578/2578), done.
remote: Total 2978 (delta 507), reused 1794 (delta 324), pack-reused 0
Receiving objects: 100% (2978/2978), 376.92 MiB | 19.63 MiB/s, done.
Resolving deltas: 100% (507/507), 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
2018-10-09 16:17:02.716687: 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:
        loggingmetrichook
        loggingtensorhook
        examplespersecondhook
        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
2018-10-09 16:17:05.048463: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
I1009 16:17:05.053292 140061072475904 tf_logging.py:115] Using config: {'_num_ps_replicas': 0, '_eval_distribute': None, '_log_step_count_steps': 100, '_service': None, '_session_config': device_count {
  key: "GPU"
}
, '_device_fn': None, '_train_distribute': None, '_protocol': None, '_save_summary_steps': 100, '_master': '', '_evaluation_master': '', '_model_dir': '/tmp/census_model', '_num_worker_replicas': 1, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f622d45bba8>, '_task_id': 0, '_tf_random_seed': None, '_global_id_in_cluster': 0, '_save_checkpoints_secs': 600, '_save_checkpoints_steps': None, '_keep_checkpoint_max': 5, '_is_chief': True, '_keep_checkpoint_every_n_hours': 10000, '_experimental_distribute': None, '_task_type': 'worker'}
W1009 16:17:05.054460 140061072475904 tf_logging.py:120] 'cpuinfo' not imported. CPU info will not be logged.
W1009 16:17:05.054889 140061072475904 tf_logging.py:120] 'psutil' not imported. Memory info will not be logged.
I1009 16:17:05.192945 140061072475904 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'}], 'run_date': '2018-10-09T16:17:05.053988Z', 'model_name': 'wide_deep', 'test_id': None, 'tensorflow_version': {'git_hash': "b'unknown'", 'version': '1.11.0'}, 'machine_config': {'gpu_info': {'count': 0}}, 'tensorflow_environment_variables': [], 'dataset': {'name': 'Census Income'}}
I1009 16:17:05.237414 140061072475904 tf_logging.py:115] Parsing /tmp/census_data/adult.data
I1009 16:17:05.273519 140061072475904 tf_logging.py:115] Calling model_fn.
I1009 16:17:06.173745 140061072475904 tf_logging.py:115] Done calling model_fn.
I1009 16:17:06.174018 140061072475904 tf_logging.py:115] Create CheckpointSaverHook.
I1009 16:17:06.549936 140061072475904 tf_logging.py:115] Graph was finalized.
I1009 16:17:06.652294 140061072475904 tf_logging.py:115] Running local_init_op.
I1009 16:17:06.678956 140061072475904 tf_logging.py:115] Done running local_init_op.
I1009 16:17:07.380616 140061072475904 tf_logging.py:115] Saving checkpoints for 0 into /tmp/census_model/model.ckpt.
I1009 16:17:07.949594 140061072475904 tf_logging.py:115] average_loss = 0.6931472, loss = 27.725887
I1009 16:17:07.950152 140061072475904 tf_logging.py:115] loss = 27.725887, step = 1
I1009 16:17:08.538611 140061072475904 tf_logging.py:115] global_step/sec: 169.622
I1009 16:17:08.539464 140061072475904 tf_logging.py:115] average_loss = 0.39386895, loss = 15.754758 (0.590 sec)
I1009 16:17:08.539777 140061072475904 tf_logging.py:115] loss = 15.754758, step = 101 (0.590 sec)
I1009 16:17:08.878756 140061072475904 tf_logging.py:115] global_step/sec: 293.969
I1009 16:17:08.879661 140061072475904 tf_logging.py:115] average_loss = 0.34996048, loss = 13.998419 (0.340 sec)
I1009 16:17:08.879964 140061072475904 tf_logging.py:115] loss = 13.998419, step = 201 (0.340 sec)
I1009 16:17:09.214598 140061072475904 tf_logging.py:115] global_step/sec: 297.757
I1009 16:17:09.215559 140061072475904 tf_logging.py:115] average_loss = 0.23998952, loss = 9.599581 (0.336 sec)
I1009 16:17:09.215908 140061072475904 tf_logging.py:115] loss = 9.599581, step = 301 (0.336 sec)
I1009 16:17:09.553291 140061072475904 tf_logging.py:115] global_step/sec: 295.249
I1009 16:17:09.554064 140061072475904 tf_logging.py:115] average_loss = 0.3641164, loss = 14.564656 (0.339 sec)
I1009 16:17:09.554347 140061072475904 tf_logging.py:115] loss = 14.564656, step = 401 (0.338 sec)
I1009 16:17:09.888134 140061072475904 tf_logging.py:115] global_step/sec: 298.656
I1009 16:17:09.888899 140061072475904 tf_logging.py:115] average_loss = 0.37377656, loss = 14.951062 (0.335 sec)
I1009 16:17:09.889189 140061072475904 tf_logging.py:115] loss = 14.951062, step = 501 (0.335 sec)
I1009 16:17:10.230372 140061072475904 tf_logging.py:115] global_step/sec: 292.182
I1009 16:17:10.231162 140061072475904 tf_logging.py:115] average_loss = 0.44428167, loss = 17.771267 (0.342 sec)
I1009 16:17:10.231468 140061072475904 tf_logging.py:115] loss = 17.771267, step = 601 (0.342 sec)
I1009 16:17:10.540647 140061072475904 tf_logging.py:115] global_step/sec: 322.283
I1009 16:17:10.541343 140061072475904 tf_logging.py:115] average_loss = 0.35957533, loss = 14.383014 (0.310 sec)
I1009 16:17:10.541604 140061072475904 tf_logging.py:115] loss = 14.383014, step = 701 (0.310 sec)
I1009 16:17:10.869942 140061072475904 tf_logging.py:115] global_step/sec: 303.721
I1009 16:17:10.870860 140061072475904 tf_logging.py:115] average_loss = 0.34024596, loss = 13.6098385 (0.330 sec)
I1009 16:17:10.871265 140061072475904 tf_logging.py:115] loss = 13.6098385, step = 801 (0.330 sec)
I1009 16:17:11.255781 140061072475904 tf_logging.py:115] global_step/sec: 259.158
I1009 16:17:11.256601 140061072475904 tf_logging.py:115] average_loss = 0.41650906, loss = 16.660362 (0.386 sec)
I1009 16:17:11.256903 140061072475904 tf_logging.py:115] loss = 16.660362, step = 901 (0.386 sec)
I1009 16:17:11.594015 140061072475904 tf_logging.py:115] global_step/sec: 295.654
I1009 16:17:11.594866 140061072475904 tf_logging.py:115] average_loss = 0.332558, loss = 13.3023205 (0.338 sec)
I1009 16:17:11.595208 140061072475904 tf_logging.py:115] loss = 13.3023205, step = 1001 (0.338 sec)
I1009 16:17:11.937734 140061072475904 tf_logging.py:115] global_step/sec: 290.926
I1009 16:17:11.938567 140061072475904 tf_logging.py:115] average_loss = 0.2967804, loss = 11.871216 (0.344 sec)
I1009 16:17:11.938887 140061072475904 tf_logging.py:115] loss = 11.871216, step = 1101 (0.344 sec)
I1009 16:17:12.280265 140061072475904 tf_logging.py:115] global_step/sec: 291.968
I1009 16:17:12.281141 140061072475904 tf_logging.py:115] average_loss = 0.40725145, loss = 16.290058 (0.343 sec)
I1009 16:17:12.281456 140061072475904 tf_logging.py:115] loss = 16.290058, step = 1201 (0.343 sec)
I1009 16:17:12.610941 140061072475904 tf_logging.py:115] global_step/sec: 302.411
I1009 16:17:12.611870 140061072475904 tf_logging.py:115] average_loss = 0.29060557, loss = 11.624223 (0.331 sec)
I1009 16:17:12.612231 140061072475904 tf_logging.py:115] loss = 11.624223, step = 1301 (0.331 sec)
I1009 16:17:12.954074 140061072475904 tf_logging.py:115] global_step/sec: 291.414
I1009 16:17:12.954836 140061072475904 tf_logging.py:115] average_loss = 0.4619999, loss = 18.479996 (0.343 sec)
I1009 16:17:12.955148 140061072475904 tf_logging.py:115] loss = 18.479996, step = 1401 (0.343 sec)
I1009 16:17:13.289181 140061072475904 tf_logging.py:115] global_step/sec: 298.435
I1009 16:17:13.290104 140061072475904 tf_logging.py:115] average_loss = 0.34678543, loss = 13.871417 (0.335 sec)
I1009 16:17:13.290464 140061072475904 tf_logging.py:115] loss = 13.871417, step = 1501 (0.335 sec)
I1009 16:17:13.629909 140061072475904 tf_logging.py:115] global_step/sec: 293.484
I1009 16:17:13.630782 140061072475904 tf_logging.py:115] average_loss = 0.2735587, loss = 10.9423485 (0.341 sec)
I1009 16:17:13.631147 140061072475904 tf_logging.py:115] loss = 10.9423485, step = 1601 (0.341 sec)
I1009 16:17:13.723575 140061072475904 tf_logging.py:115] Saving checkpoints for 1629 into /tmp/census_model/model.ckpt.
I1009 16:17:13.844036 140061072475904 tf_logging.py:115] Loss for final step: 0.046886865.
I1009 16:17:13.854052 140061072475904 tf_logging.py:115] Parsing /tmp/census_data/adult.test
I1009 16:17:13.882764 140061072475904 tf_logging.py:115] Calling model_fn.
W1009 16:17:14.954841 140061072475904 tf_logging.py:125] Trapezoidal rule is known to produce incorrect PR-AUCs; please switch to "careful_interpolation" instead.
W1009 16:17:14.972438 140061072475904 tf_logging.py:125] Trapezoidal rule is known to produce incorrect PR-AUCs; please switch to "careful_interpolation" instead.
I1009 16:17:14.991327 140061072475904 tf_logging.py:115] Done calling model_fn.
I1009 16:17:15.009677 140061072475904 tf_logging.py:115] Starting evaluation at 2018-10-09-16:17:15
I1009 16:17:15.126896 140061072475904 tf_logging.py:115] Graph was finalized.
I1009 16:17:15.128459 140061072475904 tf_logging.py:115] Restoring parameters from /tmp/census_model/model.ckpt-1629
I1009 16:17:15.175018 140061072475904 tf_logging.py:115] Running local_init_op.
I1009 16:17:15.204138 140061072475904 tf_logging.py:115] Done running local_init_op.
I1009 16:17:17.024853 140061072475904 tf_logging.py:115] Finished evaluation at 2018-10-09-16:17:17
I1009 16:17:17.025108 140061072475904 tf_logging.py:115] Saving dict for global step 1629: accuracy = 0.83711076, accuracy_baseline = 0.76377374, auc = 0.8841798, auc_precision_recall = 0.69626415, average_loss = 0.3507863, global_step = 1629, label/mean = 0.23622628, loss = 13.997921, precision = 0.6953534, prediction/mean = 0.23868126, recall = 0.5525221
I1009 16:17:17.249652 140061072475904 tf_logging.py:115] Saving 'checkpoint_path' summary for global step 1629: /tmp/census_model/model.ckpt-1629
I1009 16:17:17.250328 140061072475904 tf_logging.py:115] Results at epoch 2 / 2
I1009 16:17:17.250420 140061072475904 tf_logging.py:115] ------------------------------------------------------------
I1009 16:17:17.250502 140061072475904 tf_logging.py:115] accuracy: 0.83711076
I1009 16:17:17.250572 140061072475904 tf_logging.py:115] accuracy_baseline: 0.76377374
I1009 16:17:17.250638 140061072475904 tf_logging.py:115] auc: 0.8841798
I1009 16:17:17.250701 140061072475904 tf_logging.py:115] auc_precision_recall: 0.69626415
I1009 16:17:17.250767 140061072475904 tf_logging.py:115] average_loss: 0.3507863
I1009 16:17:17.250836 140061072475904 tf_logging.py:115] global_step: 1629
I1009 16:17:17.250901 140061072475904 tf_logging.py:115] label/mean: 0.23622628
I1009 16:17:17.250980 140061072475904 tf_logging.py:115] loss: 13.997921
I1009 16:17:17.251052 140061072475904 tf_logging.py:115] precision: 0.6953534
I1009 16:17:17.251115 140061072475904 tf_logging.py:115] prediction/mean: 0.23868126
I1009 16:17:17.251177 140061072475904 tf_logging.py:115] recall: 0.5525221
I1009 16:17:17.251313 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.8371107578277588, 'timestamp': '2018-10-09T16:17:17.251276Z', 'global_step': 1629, 'name': 'accuracy', 'extras': []}
I1009 16:17:17.251432 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.7637737393379211, 'timestamp': '2018-10-09T16:17:17.251410Z', 'global_step': 1629, 'name': 'accuracy_baseline', 'extras': []}
I1009 16:17:17.251533 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.8841797709465027, 'timestamp': '2018-10-09T16:17:17.251514Z', 'global_step': 1629, 'name': 'auc', 'extras': []}
I1009 16:17:17.251632 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.6962641477584839, 'timestamp': '2018-10-09T16:17:17.251614Z', 'global_step': 1629, 'name': 'auc_precision_recall', 'extras': []}
I1009 16:17:17.251730 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.3507862985134125, 'timestamp': '2018-10-09T16:17:17.251711Z', 'global_step': 1629, 'name': 'average_loss', 'extras': []}
I1009 16:17:17.251827 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.23622627556324005, 'timestamp': '2018-10-09T16:17:17.251809Z', 'global_step': 1629, 'name': 'label/mean', 'extras': []}
I1009 16:17:17.251924 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 13.997920989990234, 'timestamp': '2018-10-09T16:17:17.251905Z', 'global_step': 1629, 'name': 'loss', 'extras': []}
I1009 16:17:17.252021 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.6953533887863159, 'timestamp': '2018-10-09T16:17:17.252002Z', 'global_step': 1629, 'name': 'precision', 'extras': []}
I1009 16:17:17.252116 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.23868125677108765, 'timestamp': '2018-10-09T16:17:17.252098Z', 'global_step': 1629, 'name': 'prediction/mean', 'extras': []}
I1009 16:17:17.252212 140061072475904 tf_logging.py:115] Benchmark metric: {'unit': None, 'value': 0.5525221228599548, 'timestamp': '2018-10-09T16:17:17.252194Z', 'global_step': 1629, 'name': 'recall', 'extras': []}

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: ['hours_per_week', 'native_country', 'capital_loss', 'education_num', 'occupation']

A batch of Ages  : tf.Tensor([19 43 45 25 35 24 21 34 32 37], 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.
I1009 16:17:19.018875 140156615698176 tf_logging.py:115] Parsing /tmp/census_data/adult.data
Feature keys: ['hours_per_week', 'native_country', 'capital_loss', 'education_num', 'occupation']

Age batch   : tf.Tensor([37 42 27 19 29 59 46 54 52 35], shape=(10,), dtype=int32)

Label batch : tf.Tensor([False False False False False False False False  True 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([[37.],
       [42.],
       [27.],
       [19.],
       [29.],
       [59.],
       [46.],
       [54.],
       [52.],
       [35.]], 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)
{'auc': 0.6781993, 'auc_precision_recall': 0.31136066, 'average_loss': 0.52740055, 'accuracy': 0.76377374, 'label/mean': 0.23622628, 'recall': 0.0, 'precision': 0.0, 'accuracy_baseline': 0.76377374, 'global_step': 1018, 'prediction/mean': 0.20646697, 'loss': 33.672974}

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([[37.,  0.,  0., 16., 35.],
       [42.,  0.,  0.,  9., 40.],
       [27.,  0.,  0., 13., 50.],
       [19.,  0.,  0.,  9., 30.],
       [29.,  0.,  0., 13., 40.],
       [59.,  0.,  0.,  9., 35.],
       [46.,  0.,  0., 15., 48.],
       [54.,  0.,  0.,  6., 60.],
       [52.,  0.,  0.,  9., 52.],
       [35.,  0.,  0.,  9., 40.]], 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.78250724
accuracy_baseline: 0.76377374
auc: 0.70410156
auc_precision_recall: 0.49718997
average_loss: 2.7589335
global_step: 1018
label/mean: 0.23622628
loss: 176.14978
precision: 0.61526835
prediction/mean: 0.26403713
recall: 0.21164846

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

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
Machine-op-inspct
Prof-specialty
Prof-specialty
Prof-specialty
Farming-fishing
Farming-fishing
Machine-op-inspct
Protective-serv
?

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, 911, 979, 979, 979, 936, 936, 911, 684,  65])

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.79823107
accuracy_baseline: 0.76377374
auc: 0.6876817
auc_precision_recall: 0.50509375
average_loss: 3.3195345
global_step: 1018
label/mean: 0.23622628
loss: 211.94252
precision: 0.6521975
prediction/mean: 0.14528999
recall: 0.3125325

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([[37.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.],
       [42.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.],
       [27.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [19.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [29.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [59.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.],
       [46.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.],
       [54.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [52.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [35.,  0.,  0.,  0.,  0.,  1.,  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.
I1009 16:17:43.038667 140156615698176 tf_logging.py:115] Using default config.
INFO:tensorflow:Using config: {'_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_save_checkpoints_steps': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f7802f8a198>, '_eval_distribute': None, '_save_checkpoints_secs': 600, '_save_summary_steps': 100, '_global_id_in_cluster': 0, '_service': None, '_is_chief': True, '_tf_random_seed': None, '_protocol': None, '_model_dir': '/tmp/tmpews_wh40', '_task_id': 0, '_device_fn': None, '_train_distribute': None, '_num_worker_replicas': 1, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_task_type': 'worker', '_num_ps_replicas': 0, '_master': '', '_experimental_distribute': None, '_keep_checkpoint_max': 5, '_evaluation_master': ''}
I1009 16:17:43.042024 140156615698176 tf_logging.py:115] Using config: {'_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_save_checkpoints_steps': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f7802f8a198>, '_eval_distribute': None, '_save_checkpoints_secs': 600, '_save_summary_steps': 100, '_global_id_in_cluster': 0, '_service': None, '_is_chief': True, '_tf_random_seed': None, '_protocol': None, '_model_dir': '/tmp/tmpews_wh40', '_task_id': 0, '_device_fn': None, '_train_distribute': None, '_num_worker_replicas': 1, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_task_type': 'worker', '_num_ps_replicas': 0, '_master': '', '_experimental_distribute': None, '_keep_checkpoint_max': 5, '_evaluation_master': ''}

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.80
accuracy_baseline: 0.76
auc: 0.69
auc_precision_recall: 0.51
average_loss: 3.32
global_step: 1018.00
label/mean: 0.24
loss: 211.94
precision: 0.65
prediction/mean: 0.15
recall: 0.31

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.55
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.68
prediction/mean: 0.24
recall: 0.57

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.