EVOLUTION-MANAGER

Edit File: linear_testing_utils.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utils for testing linear estimators."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import os
import shutil
import tempfile

import numpy as np
import six
import tensorflow as tf
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import ops
from tensorflow.python.keras.optimizer_v2 import gradient_descent
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow_estimator.python.estimator import estimator
from tensorflow_estimator.python.estimator.canned import linear
from tensorflow_estimator.python.estimator.canned import metric_keys
from tensorflow_estimator.python.estimator.export import export
from tensorflow_estimator.python.estimator.inputs import numpy_io
from tensorflow_estimator.python.estimator.inputs import pandas_io

try:
  # pylint: disable=g-import-not-at-top
  import pandas as pd
  HAS_PANDAS = True
except IOError:
  # Pandas writes a temporary file during import. If it fails, don't use pandas.
  HAS_PANDAS = False
except ImportError:
  HAS_PANDAS = False

# pylint rules which are disabled by default for test files.
# pylint: disable=invalid-name,protected-access,missing-docstring

# Names of variables created by model.
AGE_WEIGHT_NAME = 'linear/linear_model/age/weights'
HEIGHT_WEIGHT_NAME = 'linear/linear_model/height/weights'
OCCUPATION_WEIGHT_NAME = 'linear/linear_model/occupation/weights'
BIAS_NAME = 'linear/linear_model/bias_weights'
LANGUAGE_WEIGHT_NAME = 'linear/linear_model/language/weights'

def assert_close(expected, actual, rtol=1e-04, name='assert_close'):
  with ops.name_scope(name, 'assert_close', (expected, actual, rtol)) as scope:
    expected = ops.convert_to_tensor(expected, name='expected')
    actual = ops.convert_to_tensor(actual, name='actual')
    rdiff = tf.math.abs(expected - actual, 'diff') / tf.math.abs(expected)
    rtol = ops.convert_to_tensor(rtol, name='rtol')
    return tf.compat.v1.debugging.assert_less(
        rdiff,
        rtol,
        data=('Condition expected =~ actual did not hold element-wise:'
              'expected = ', expected, 'actual = ', actual, 'rdiff = ', rdiff,
              'rtol = ', rtol,),
        name=scope)

def save_variables_to_ckpt(model_dir):
  init_all_op = [tf.compat.v1.initializers.global_variables()]
  with tf.compat.v1.Session() as sess:
    sess.run(init_all_op)
    tf.compat.v1.train.Saver().save(sess, os.path.join(model_dir, 'model.ckpt'))

def queue_parsed_features(feature_map):
  tensors_to_enqueue = []
  keys = []
  for key, tensor in six.iteritems(feature_map):
    keys.append(key)
    tensors_to_enqueue.append(tensor)
  queue_dtypes = [x.dtype for x in tensors_to_enqueue]
  input_queue = tf.queue.FIFOQueue(capacity=100, dtypes=queue_dtypes)
  tf.compat.v1.train.queue_runner.add_queue_runner(
      tf.compat.v1.train.queue_runner.QueueRunner(
          input_queue, [input_queue.enqueue(tensors_to_enqueue)]))
  dequeued_tensors = input_queue.dequeue()
  return {keys[i]: dequeued_tensors[i] for i in range(len(dequeued_tensors))}

def sorted_key_dict(unsorted_dict):
  return {k: unsorted_dict[k] for k in sorted(unsorted_dict)}

def sigmoid(x):
  return 1 / (1 + np.exp(-1.0 * x))

def mock_optimizer(testcase, expected_loss=None):
  expected_var_names = ['%s:0' % AGE_WEIGHT_NAME, '%s:0' % BIAS_NAME]

class _Optimizer(optimizer_v2.OptimizerV2):

def get_updates(self, loss, params):
      trainable_vars = params
      testcase.assertItemsEqual(expected_var_names,
                                [var.name for var in trainable_vars])

# Verify loss. We can't check the value directly, so we add an assert op.
      testcase.assertEquals(0, loss.shape.ndims)
      if expected_loss is None:
        if self.iterations is not None:
          return [self.iterations.assign_add(1).op]
        return [tf.no_op()]
      assert_loss = assert_close(
          tf.cast(expected_loss, name='expected', dtype=tf.dtypes.float32),
          loss,
          name='assert_loss')
      with tf.control_dependencies((assert_loss,)):
        if self.iterations is not None:
          return [self.iterations.assign_add(1).op]
        return [tf.no_op()]

def get_config(self):
      config = super(_Optimizer, self).get_config()
      return config

optimizer = _Optimizer(name='my_optimizer')

return optimizer

# TODO(b/36813849): Add tests with dynamic shape inputs using placeholders.
class BaseLinearRegressorEvaluationTest(object):

def __init__(self, linear_regressor_fn, fc_lib=feature_column_v2):
    self._linear_regressor_fn = linear_regressor_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

def test_evaluation_for_simple_data(self):
    with tf.Graph().as_default():
      tf.Variable([[11.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([2.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        model_dir=self._model_dir)
    eval_metrics = linear_regressor.evaluate(
        input_fn=lambda: ({
            'age': ((1,),)
        }, ((10.,),)), steps=1)

# Logit is (1. * 11.0 + 2.0) = 13, while label is 10. Loss is 3**2 = 9.
    self.assertDictEqual(
        {
            metric_keys.MetricKeys.LOSS: 9.,
            metric_keys.MetricKeys.LOSS_MEAN: 9.,
            metric_keys.MetricKeys.PREDICTION_MEAN: 13.,
            metric_keys.MetricKeys.LABEL_MEAN: 10.,
            tf.compat.v1.GraphKeys.GLOBAL_STEP: 100
        }, eval_metrics)

def test_evaluation_batch(self):
    """Tests evaluation for batch_size==2."""
    with tf.Graph().as_default():
      tf.Variable([[11.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([2.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# Logit is (1. * 11.0 + 2.0) = 13, while label is 10.
    # Loss per example is 3**2 = 9.
    # Training loss is the sum over batch size = (9 + 9) / 2 = 9
    # Average loss is the average over batch = 9
    self.assertDictEqual(
        {
            metric_keys.MetricKeys.LOSS: 9.,
            metric_keys.MetricKeys.LOSS_MEAN: 9.,
            metric_keys.MetricKeys.PREDICTION_MEAN: 13.,
            metric_keys.MetricKeys.LABEL_MEAN: 10.,
            tf.compat.v1.GraphKeys.GLOBAL_STEP: 100
        }, eval_metrics)

def test_evaluation_weights(self):
    """Tests evaluation with weights."""
    with tf.Graph().as_default():
      tf.Variable([[11.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([2.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

def _input_fn():
      features = {'age': ((1,), (1,)), 'weights': ((1.,), (2.,))}
      labels = ((10.,), (10.,))
      return features, labels

linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        weight_column='weights',
        model_dir=self._model_dir)
    eval_metrics = linear_regressor.evaluate(input_fn=_input_fn, steps=1)

# Logit is (1. * 11.0 + 2.0) = 13, while label is 10.
    # Loss per example is 3**2 = 9.
    # Training loss is the weighted sum over batch / batch size =
    #     (9 + 2*9) / 2 = 13.5
    # average loss is the weighted average = 9 + 2*9 / (1 + 2) = 9
    self.assertDictEqual(
        {
            metric_keys.MetricKeys.LOSS: 13.5,
            metric_keys.MetricKeys.LOSS_MEAN: 9.,
            metric_keys.MetricKeys.PREDICTION_MEAN: 13.,
            metric_keys.MetricKeys.LABEL_MEAN: 10.,
            tf.compat.v1.GraphKeys.GLOBAL_STEP: 100
        }, eval_metrics)

def test_evaluation_for_multi_dimensions(self):
    x_dim = 3
    label_dim = 2
    with tf.Graph().as_default():
      tf.Variable([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([7.0, 8.0], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age', shape=(x_dim,)),),
        label_dimension=label_dim,
        model_dir=self._model_dir)
    input_fn = numpy_io.numpy_input_fn(
        x={
            'age': np.array([[2., 4., 5.]]),
        },
        y=np.array([[46., 58.]]),
        batch_size=1,
        num_epochs=None,
        shuffle=False)
    eval_metrics = linear_regressor.evaluate(input_fn=input_fn, steps=1)

self.assertItemsEqual(
        (metric_keys.MetricKeys.LOSS, metric_keys.MetricKeys.LOSS_MEAN,
         metric_keys.MetricKeys.PREDICTION_MEAN,
         metric_keys.MetricKeys.LABEL_MEAN, tf.compat.v1.GraphKeys.GLOBAL_STEP),
        eval_metrics.keys())

# Logit is
    #   [2., 4., 5.] * [1.0, 2.0] + [7.0, 8.0] = [39, 50] + [7.0, 8.0]
    #                  [3.0, 4.0]
    #                  [5.0, 6.0]
    # which is [46, 58]
    self.assertAlmostEqual(0, eval_metrics[metric_keys.MetricKeys.LOSS])

def test_evaluation_for_multiple_feature_columns(self):
    with tf.Graph().as_default():
      tf.Variable([[10.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([[2.0]], name=HEIGHT_WEIGHT_NAME)
      tf.Variable([5.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

batch_size = 2
    feature_columns = [
        self._fc_lib.numeric_column('age'),
        self._fc_lib.numeric_column('height')
    ]
    input_fn = numpy_io.numpy_input_fn(
        x={
            'age': np.array([20, 40]),
            'height': np.array([4, 8])
        },
        y=np.array([[213.], [421.]]),
        batch_size=batch_size,
        num_epochs=None,
        shuffle=False)

est = self._linear_regressor_fn(
        feature_columns=feature_columns, model_dir=self._model_dir)

eval_metrics = est.evaluate(input_fn=input_fn, steps=1)
    self.assertItemsEqual(
        (metric_keys.MetricKeys.LOSS, metric_keys.MetricKeys.LOSS_MEAN,
         metric_keys.MetricKeys.PREDICTION_MEAN,
         metric_keys.MetricKeys.LABEL_MEAN, tf.compat.v1.GraphKeys.GLOBAL_STEP),
        eval_metrics.keys())

# Logit is [(20. * 10.0 + 4 * 2.0 + 5.0), (40. * 10.0 + 8 * 2.0 + 5.0)] =
    # [213.0, 421.0], while label is [213., 421.]. Loss = 0.
    self.assertAlmostEqual(0, eval_metrics[metric_keys.MetricKeys.LOSS])

def test_evaluation_for_multiple_feature_columns_mix(self):
    with tf.Graph().as_default():
      tf.Variable([[10.0]], name=AGE_WEIGHT_NAME)
      tf.Variable([[2.0]], name=HEIGHT_WEIGHT_NAME)
      tf.Variable([5.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

batch_size = 2
    feature_columns = [
        tf.feature_column.numeric_column('age'),
        tf.feature_column.numeric_column('height')
    ]

def _input_fn():
      features_ds = tf.compat.v1.data.Dataset.from_tensor_slices({
          'age': np.array([20, 40]),
          'height': np.array([4, 8])
      })
      labels_ds = tf.compat.v1.data.Dataset.from_tensor_slices(
          np.array([[213.], [421.]]))
      return (tf.compat.v1.data.Dataset.zip(
          (features_ds, labels_ds)).batch(batch_size).repeat(None))

est = self._linear_regressor_fn(
        feature_columns=feature_columns, model_dir=self._model_dir)

eval_metrics = est.evaluate(input_fn=_input_fn, steps=1)
    self.assertItemsEqual(
        (metric_keys.MetricKeys.LOSS, metric_keys.MetricKeys.LOSS_MEAN,
         metric_keys.MetricKeys.PREDICTION_MEAN,
         metric_keys.MetricKeys.LABEL_MEAN, tf.compat.v1.GraphKeys.GLOBAL_STEP),
        eval_metrics.keys())

class BaseLinearRegressorPredictTest(object):

def __init__(self, linear_regressor_fn, fc_lib=feature_column_v2):
    self._linear_regressor_fn = linear_regressor_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

def test_1d(self):
    """Tests predict when all variables are one-dimensional."""
    with tf.Graph().as_default():
      tf.Variable([[10.]], name='linear/linear_model/x/weights')
      tf.Variable([.2], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('x'),),
        model_dir=self._model_dir)

predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[2.]])},
        y=None,
        batch_size=1,
        num_epochs=1,
        shuffle=False)
    predictions = linear_regressor.predict(input_fn=predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # x * weight + bias = 2. * 10. + .2 = 20.2
    self.assertAllClose([[20.2]], predicted_scores)

def testMultiDim(self):
    """Tests predict when all variables are multi-dimenstional."""
    batch_size = 2
    label_dimension = 3
    x_dim = 4
    feature_columns = (self._fc_lib.numeric_column('x', shape=(x_dim,)),)
    with tf.Graph().as_default():
      tf.Variable(  # shape=[x_dim, label_dimension]
          [[1., 2., 3.], [2., 3., 4.], [3., 4., 5.], [4., 5., 6.]],
          name='linear/linear_model/x/weights')
      tf.Variable(  # shape=[label_dimension]
          [.2, .4, .6], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)

predict_input_fn = numpy_io.numpy_input_fn(
        # x shape=[batch_size, x_dim]
        x={'x': np.array([[1., 2., 3., 4.], [5., 6., 7., 8.]])},
        y=None,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)
    predictions = linear_regressor.predict(input_fn=predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # score = x * weight + bias, shape=[batch_size, label_dimension]
    self.assertAllClose([[30.2, 40.4, 50.6], [70.2, 96.4, 122.6]],
                        predicted_scores)

def testTwoFeatureColumns(self):
    """Tests predict with two feature columns."""
    with tf.Graph().as_default():
      tf.Variable([[10.]], name='linear/linear_model/x0/weights')
      tf.Variable([[20.]], name='linear/linear_model/x1/weights')
      tf.Variable([.2], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('x0'),
                         self._fc_lib.numeric_column('x1')),
        model_dir=self._model_dir)

predict_input_fn = numpy_io.numpy_input_fn(
        x={
            'x0': np.array([[2.]]),
            'x1': np.array([[3.]])
        },
        y=None,
        batch_size=1,
        num_epochs=1,
        shuffle=False)
    predictions = linear_regressor.predict(input_fn=predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # x0 * weight0 + x1 * weight1 + bias = 2. * 10. + 3. * 20 + .2 = 80.2
    self.assertAllClose([[80.2]], predicted_scores)

def testTwoFeatureColumnsMix(self):
    """Tests predict with two feature columns."""
    with tf.Graph().as_default():
      tf.Variable([[10.]], name='linear/linear_model/x0/weights')
      tf.Variable([[20.]], name='linear/linear_model/x1/weights')
      tf.Variable([.2], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

linear_regressor = self._linear_regressor_fn(
        feature_columns=(tf.feature_column.numeric_column('x0'),
                         tf.feature_column.numeric_column('x1')),
        model_dir=self._model_dir)

def _predict_input_fn():
      return tf.compat.v1.data.Dataset.from_tensor_slices({
          'x0': np.array([[2.]]),
          'x1': np.array([[3.]])
      }).batch(1)

predictions = linear_regressor.predict(input_fn=_predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # x0 * weight0 + x1 * weight1 + bias = 2. * 10. + 3. * 20 + .2 = 80.2
    self.assertAllClose([[80.2]], predicted_scores)

def testSparseCombiner(self):
    w_a = 2.0
    w_b = 3.0
    w_c = 5.0
    bias = 5.0
    with tf.Graph().as_default():
      tf.Variable([[w_a], [w_b], [w_c]], name=LANGUAGE_WEIGHT_NAME)
      tf.Variable([bias], name=BIAS_NAME)
      tf.Variable(
          1, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

def _input_fn():
      return tf.compat.v1.data.Dataset.from_tensors({
          'language':
              tf.sparse.SparseTensor(
                  values=['a', 'c', 'b', 'c'],
                  indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
                  dense_shape=[2, 2]),
      })

feature_columns = (self._fc_lib.categorical_column_with_vocabulary_list(
        'language', vocabulary_list=['a', 'b', 'c']),)

# Check prediction for each sparse_combiner.
    # With sparse_combiner = 'sum', we have
    # logits_1 = w_a + w_c + bias
    #          = 2.0 + 5.0 + 5.0 = 12.0
    # logits_2 = w_b + w_c + bias
    #          = 3.0 + 5.0 + 5.0 = 13.0
    linear_regressor = self._linear_regressor_fn(
        feature_columns=feature_columns, model_dir=self._model_dir)
    predictions = linear_regressor.predict(input_fn=_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    self.assertAllClose([[12.0], [13.0]], predicted_scores)

# With sparse_combiner = 'mean', we have
    # logits_1 = 1/2 * (w_a + w_c) + bias
    #          = 1/2 * (2.0 + 5.0) + 5.0 = 8.5
    # logits_2 = 1/2 * (w_b + w_c) + bias
    #          = 1/2 * (3.0 + 5.0) + 5.0 = 9.0
    linear_regressor = self._linear_regressor_fn(
        feature_columns=feature_columns,
        model_dir=self._model_dir,
        sparse_combiner='mean')
    predictions = linear_regressor.predict(input_fn=_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    self.assertAllClose([[8.5], [9.0]], predicted_scores)

# With sparse_combiner = 'sqrtn', we have
    # logits_1 = sqrt(2)/2 * (w_a + w_c) + bias
    #          = sqrt(2)/2 * (2.0 + 5.0) + 5.0 = 9.94974
    # logits_2 = sqrt(2)/2 * (w_b + w_c) + bias
    #          = sqrt(2)/2 * (3.0 + 5.0) + 5.0 = 10.65685
    linear_regressor = self._linear_regressor_fn(
        feature_columns=feature_columns,
        model_dir=self._model_dir,
        sparse_combiner='sqrtn')
    predictions = linear_regressor.predict(input_fn=_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    self.assertAllClose([[9.94974], [10.65685]], predicted_scores)

class BaseLinearRegressorIntegrationTest(object):

def __init__(self, linear_regressor_fn, fc_lib=feature_column_v2):
    self._linear_regressor_fn = linear_regressor_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                          input_dimension, label_dimension, prediction_length):
    feature_columns = [
        self._fc_lib.numeric_column('x', shape=(input_dimension,))
    ]
    est = self._linear_regressor_fn(
        feature_columns=feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)

# TRAIN
    # learn y = x
    est.train(train_input_fn, steps=200)

# EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(200, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn(metric_keys.MetricKeys.LOSS, six.iterkeys(scores))

# PREDICT
    predictions = np.array(
        [x['predictions'] for x in est.predict(predict_input_fn)])
    self.assertAllEqual((prediction_length, label_dimension), predictions.shape)

# EXPORT
    feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                        serving_input_receiver_fn)
    self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))

def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    label_dimension = 2
    input_dimension = label_dimension
    batch_size = 10
    prediction_length = batch_size
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=None,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)

self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        label_dimension=label_dimension,
        prediction_length=prediction_length)

def test_pandas_input_fn(self):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return

# Pandas DataFrame natually supports 1 dim data only.
    label_dimension = 1
    input_dimension = label_dimension
    batch_size = 10
    data = np.array([1., 2., 3., 4.], dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(data)
    prediction_length = 4

train_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, num_epochs=None, shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x, batch_size=batch_size, shuffle=False)

def test_input_fn_from_parse_example(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    label_dimension = 2
    input_dimension = label_dimension
    batch_size = 10
    prediction_length = batch_size
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

serialized_examples = []
    for datum in data:
      example = example_pb2.Example(
          features=feature_pb2.Features(
              feature={
                  'x':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(value=datum)),
                  'y':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(
                              value=datum[:label_dimension])),
              }))
      serialized_examples.append(example.SerializeToString())

feature_spec = {
        'x': tf.io.FixedLenFeature([input_dimension], tf.dtypes.float32),
        'y': tf.io.FixedLenFeature([label_dimension], tf.dtypes.float32),
    }

def _train_input_fn():
      feature_map = tf.compat.v1.io.parse_example(serialized_examples,
                                                  feature_spec)
      features = queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

def _eval_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

def _predict_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=input_dimension,
        label_dimension=label_dimension,
        prediction_length=prediction_length)

class BaseLinearRegressorTrainingTest(object):

def __init__(self, linear_regressor_fn, fc_lib=feature_column_v2):
    self._linear_regressor_fn = linear_regressor_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

def _assert_checkpoint(self,
                         expected_global_step,
                         expected_age_weight=None,
                         expected_bias=None):
    shapes = {
        name: shape
        for (name, shape) in tf.train.list_variables(self._model_dir)
    }

self.assertEqual([], shapes[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertEqual(
        expected_global_step,
        tf.train.load_variable(self._model_dir,
                               tf.compat.v1.GraphKeys.GLOBAL_STEP))

self.assertEqual([1, 1], shapes[AGE_WEIGHT_NAME])
    if expected_age_weight is not None:
      self.assertEqual(expected_age_weight,
                       tf.train.load_variable(self._model_dir, AGE_WEIGHT_NAME))

self.assertEqual([1], shapes[BIAS_NAME])
    if expected_bias is not None:
      self.assertEqual(expected_bias,
                       tf.train.load_variable(self._model_dir, BIAS_NAME))

def testFromScratchWithDefaultOptimizer(self):
    # Create LinearRegressor.
    label = 5.
    age = 17
    linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        model_dir=self._model_dir)

# Train for a few steps, and validate final checkpoint.
    num_steps = 10
    linear_regressor.train(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=num_steps)
    self._assert_checkpoint(num_steps)

def testTrainWithOneDimLabel(self):
    label_dimension = 1
    batch_size = 20
    feature_columns = [self._fc_lib.numeric_column('age', shape=(1,))]
    est = self._linear_regressor_fn(
        feature_columns=feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)
    data_rank_1 = np.linspace(0., 2., batch_size, dtype=np.float32)
    self.assertEqual((batch_size,), data_rank_1.shape)

train_input_fn = numpy_io.numpy_input_fn(
        x={'age': data_rank_1},
        y=data_rank_1,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    est.train(train_input_fn, steps=200)
    self._assert_checkpoint(200)

def testTrainWithOneDimWeight(self):
    label_dimension = 1
    batch_size = 20
    feature_columns = [self._fc_lib.numeric_column('age', shape=(1,))]
    est = self._linear_regressor_fn(
        feature_columns=feature_columns,
        label_dimension=label_dimension,
        weight_column='w',
        model_dir=self._model_dir)

data_rank_1 = np.linspace(0., 2., batch_size, dtype=np.float32)
    self.assertEqual((batch_size,), data_rank_1.shape)

train_input_fn = numpy_io.numpy_input_fn(
        x={
            'age': data_rank_1,
            'w': data_rank_1
        },
        y=data_rank_1,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    est.train(train_input_fn, steps=200)
    self._assert_checkpoint(200)

def testFromScratch(self):
    # Create LinearRegressor.
    label = 5.
    age = 17
    # loss = (logits - label)^2 = (0 - 5.)^2 = 25.
    mock_opt = mock_optimizer(self, expected_loss=25.)
    linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        model_dir=self._model_dir,
        optimizer=mock_opt)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    linear_regressor.train(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=num_steps)
    self.assertEqual(
        num_steps,
        linear_regressor.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        expected_global_step=num_steps,
        expected_age_weight=0.,
        expected_bias=0.)

def testFromCheckpoint(self):
    # Create initial checkpoint.
    age_weight = 10.0
    bias = 5.0
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable([[age_weight]], name=AGE_WEIGHT_NAME)
      tf.Variable([bias], name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# logits = age * age_weight + bias = 17 * 10. + 5. = 175
    # loss = (logits - label)^2 = (175 - 5)^2 = 28900
    mock_opt = mock_optimizer(self, expected_loss=28900.)
    linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        model_dir=self._model_dir,
        optimizer=mock_opt)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    linear_regressor.train(
        input_fn=lambda: ({
            'age': ((17,),)
        }, ((5.,),)), steps=num_steps)
    self.assertEqual(
        initial_global_step + num_steps,
        linear_regressor.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        expected_global_step=initial_global_step + num_steps,
        expected_age_weight=age_weight,
        expected_bias=bias)

def testFromCheckpointMultiBatch(self):
    # Create initial checkpoint.
    age_weight = 10.0
    bias = 5.0
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable([[age_weight]], name=AGE_WEIGHT_NAME)
      tf.Variable([bias], name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# logits = age * age_weight + bias
    # logits[0] = 17 * 10. + 5. = 175
    # logits[1] = 15 * 10. + 5. = 155
    # loss = sum(logits - label)^2 = (175 - 5)^2 + (155 - 3)^2 = 52004
    # expected_loss = loss / 2 = 26002
    mock_opt = mock_optimizer(self, expected_loss=26002.)
    linear_regressor = self._linear_regressor_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        model_dir=self._model_dir,
        optimizer=mock_opt)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    linear_regressor.train(
        input_fn=lambda: ({
            'age': ((17,), (15,))
        }, ((5.,), (3.,))),
        steps=num_steps)
    self.assertEqual(
        initial_global_step + num_steps,
        linear_regressor.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        expected_global_step=initial_global_step + num_steps,
        expected_age_weight=age_weight,
        expected_bias=bias)

class BaseLinearClassifierTrainingTest(object):

def __init__(self, linear_classifier_fn, fc_lib=feature_column_v2):
    self._linear_classifier_fn = linear_classifier_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

def _assert_checkpoint(self,
                         n_classes,
                         expected_global_step,
                         expected_age_weight=None,
                         expected_bias=None):
    logits_dimension = n_classes if n_classes > 2 else 1

shapes = {
        name: shape
        for (name, shape) in tf.train.list_variables(self._model_dir)
    }

self.assertEqual([1, logits_dimension], shapes[AGE_WEIGHT_NAME])
    if expected_age_weight is not None:
      self.assertAllEqual(
          expected_age_weight,
          tf.train.load_variable(self._model_dir, AGE_WEIGHT_NAME))

self.assertEqual([logits_dimension], shapes[BIAS_NAME])
    if expected_bias is not None:
      self.assertAllEqual(expected_bias,
                          tf.train.load_variable(self._model_dir, BIAS_NAME))

def _testFromScratchWithDefaultOptimizer(self, n_classes):
    label = 0
    age = 17
    est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        model_dir=self._model_dir)

# Train for a few steps, and validate final checkpoint.
    num_steps = 10
    est.train(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=num_steps)
    self._assert_checkpoint(n_classes, num_steps)

def testBinaryClassesFromScratchWithDefaultOptimizer(self):
    self._testFromScratchWithDefaultOptimizer(n_classes=2)

def testMultiClassesFromScratchWithDefaultOptimizer(self):
    self._testFromScratchWithDefaultOptimizer(n_classes=4)

def _testTrainWithTwoDimsLabel(self, n_classes):
    batch_size = 20

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        model_dir=self._model_dir)
    data_rank_1 = np.array([0, 1])
    data_rank_2 = np.array([[0], [1]])
    self.assertEqual((2,), data_rank_1.shape)
    self.assertEqual((2, 1), data_rank_2.shape)

train_input_fn = numpy_io.numpy_input_fn(
        x={'age': data_rank_1},
        y=data_rank_2,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    est.train(train_input_fn, steps=200)
    self._assert_checkpoint(n_classes, 200)

def testBinaryClassesTrainWithTwoDimsLabel(self):
    self._testTrainWithTwoDimsLabel(n_classes=2)

def testMultiClassesTrainWithTwoDimsLabel(self):
    self._testTrainWithTwoDimsLabel(n_classes=4)

def _testTrainWithOneDimLabel(self, n_classes):
    batch_size = 20

def testBinaryClassesTrainWithOneDimLabel(self):
    self._testTrainWithOneDimLabel(n_classes=2)

def testMultiClassesTrainWithOneDimLabel(self):
    self._testTrainWithOneDimLabel(n_classes=4)

def _testTrainWithTwoDimsWeight(self, n_classes):
    batch_size = 20

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        weight_column='w',
        n_classes=n_classes,
        model_dir=self._model_dir)
    data_rank_1 = np.array([0, 1])
    data_rank_2 = np.array([[0], [1]])
    self.assertEqual((2,), data_rank_1.shape)
    self.assertEqual((2, 1), data_rank_2.shape)

train_input_fn = numpy_io.numpy_input_fn(
        x={
            'age': data_rank_1,
            'w': data_rank_2
        },
        y=data_rank_1,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    est.train(train_input_fn, steps=200)
    self._assert_checkpoint(n_classes, 200)

def testBinaryClassesTrainWithTwoDimsWeight(self):
    self._testTrainWithTwoDimsWeight(n_classes=2)

def testMultiClassesTrainWithTwoDimsWeight(self):
    self._testTrainWithTwoDimsWeight(n_classes=4)

def _testTrainWithOneDimWeight(self, n_classes):
    batch_size = 20

def testBinaryClassesTrainWithOneDimWeight(self):
    self._testTrainWithOneDimWeight(n_classes=2)

def testMultiClassesTrainWithOneDimWeight(self):
    self._testTrainWithOneDimWeight(n_classes=4)

def _testFromScratch(self, n_classes):
    label = 1
    age = 17
    # For binary classifier:
    #   loss = sigmoid_cross_entropy(logits, label) where logits=0 (weights are
    #   all zero initially) and label = 1 so,
    #      loss = 1 * -log ( sigmoid(logits) ) = 0.69315
    # For multi class classifier:
    #   loss = cross_entropy(logits, label) where logits are all 0s (weights are
    #   all zero initially) and label = 1 so,
    #      loss = 1 * -log ( 1.0 / n_classes )
    # For this particular test case, as logits are same, the formular
    # 1 * -log ( 1.0 / n_classes ) covers both binary and multi class cases.
    mock_opt = mock_optimizer(
        self, expected_loss=-1 * math.log(1.0 / n_classes))

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        optimizer=mock_opt,
        model_dir=self._model_dir)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    est.train(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=num_steps)
    self.assertEqual(num_steps,
                     est.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        n_classes,
        expected_global_step=num_steps,
        expected_age_weight=[[0.]] if n_classes == 2 else [[0.] * n_classes],
        expected_bias=[0.] if n_classes == 2 else [.0] * n_classes)

def testBinaryClassesFromScratch(self):
    self._testFromScratch(n_classes=2)

def testMultiClassesFromScratch(self):
    self._testFromScratch(n_classes=4)

def _testFromCheckpoint(self, n_classes):
    # Create initial checkpoint.
    label = 1
    age = 17
    # For binary case, the expected weight has shape (1,1). For multi class
    # case, the shape is (1, n_classes). In order to test the weights, set
    # weights as 2.0 * range(n_classes).
    age_weight = [[2.0]] if n_classes == 2 else (np.reshape(
        2.0 * np.array(list(range(n_classes)), dtype=np.float32),
        (1, n_classes)))
    bias = [-35.0] if n_classes == 2 else [-35.0] * n_classes
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# For binary classifier:
    #   logits = age * age_weight + bias = 17 * 2. - 35. = -1.
    #   loss = sigmoid_cross_entropy(logits, label)
    #   so, loss = 1 * -log ( sigmoid(-1) ) = 1.3133
    # For multi class classifier:
    #   loss = cross_entropy(logits, label)
    #   where logits = 17 * age_weight + bias and label = 1
    #   so, loss = 1 * -log ( soft_max(logits)[1] )
    if n_classes == 2:
      expected_loss = 1.3133
    else:
      logits = age_weight * age + bias
      logits_exp = np.exp(logits)
      softmax = logits_exp / logits_exp.sum()
      expected_loss = -1 * math.log(softmax[0, label])

mock_opt = mock_optimizer(self, expected_loss=expected_loss)

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        optimizer=mock_opt,
        model_dir=self._model_dir)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    est.train(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=num_steps)
    self.assertEqual(initial_global_step + num_steps,
                     est.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        n_classes,
        expected_global_step=initial_global_step + num_steps,
        expected_age_weight=age_weight,
        expected_bias=bias)

def testBinaryClassesFromCheckpoint(self):
    self._testFromCheckpoint(n_classes=2)

def testMultiClassesFromCheckpoint(self):
    self._testFromCheckpoint(n_classes=4)

def _testFromCheckpointFloatLabels(self, n_classes):
    """Tests float labels for binary classification."""
    # Create initial checkpoint.
    if n_classes > 2:
      return
    label = 0.8
    age = 17
    age_weight = [[2.0]]
    bias = [-35.0]
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# logits = age * age_weight + bias = 17 * 2. - 35. = -1.
    # loss = sigmoid_cross_entropy(logits, label)
    # => loss = -0.8 * log(sigmoid(-1)) -0.2 * log(sigmoid(+1)) = 1.1132617
    mock_opt = mock_optimizer(self, expected_loss=1.1132617)

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        optimizer=mock_opt,
        model_dir=self._model_dir)

def testBinaryClassesFromCheckpointFloatLabels(self):
    self._testFromCheckpointFloatLabels(n_classes=2)

def testMultiClassesFromCheckpointFloatLabels(self):
    self._testFromCheckpointFloatLabels(n_classes=4)

def _testFromCheckpointMultiBatch(self, n_classes):
    # Create initial checkpoint.
    label = [1, 0]
    age = [17.0, 18.5]
    # For binary case, the expected weight has shape (1,1). For multi class
    # case, the shape is (1, n_classes). In order to test the weights, set
    # weights as 2.0 * range(n_classes).
    age_weight = [[2.0]] if n_classes == 2 else (np.reshape(
        2.0 * np.array(list(range(n_classes)), dtype=np.float32),
        (1, n_classes)))
    bias = [-35.0] if n_classes == 2 else [-35.0] * n_classes
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

# For binary classifier:
    #   logits = age * age_weight + bias
    #   logits[0] = 17 * 2. - 35. = -1.
    #   logits[1] = 18.5 * 2. - 35. = 2.
    #   loss = sigmoid_cross_entropy(logits, label)
    #   so, loss[0] = 1 * -log ( sigmoid(-1) ) = 1.3133
    #       loss[1] = (1 - 0) * -log ( 1- sigmoid(2) ) = 2.1269
    #   expected_loss = (loss[0] + loss[1]) / batch size (2)
    # For multi class classifier:
    #   loss = cross_entropy(logits, label)
    #   where logits = [17, 18.5] * age_weight + bias and label = [1, 0]
    #   so, loss = 1 * -log ( soft_max(logits)[label] )
    #   expected_loss = (loss[0] + loss[1]) / batch size (2)
    if n_classes == 2:
      expected_loss = (1.3133 + 2.1269) / 2
    else:
      logits = age_weight * np.reshape(age, (2, 1)) + bias
      logits_exp = np.exp(logits)
      softmax_row_0 = logits_exp[0] / logits_exp[0].sum()
      softmax_row_1 = logits_exp[1] / logits_exp[1].sum()
      expected_loss_0 = -1 * math.log(softmax_row_0[label[0]])
      expected_loss_1 = -1 * math.log(softmax_row_1[label[1]])
      expected_loss = (expected_loss_0 + expected_loss_1) / 2

mock_opt = mock_optimizer(self, expected_loss=expected_loss)

est = linear.LinearClassifierV2(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        optimizer=mock_opt,
        model_dir=self._model_dir)

# Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    est.train(input_fn=lambda: ({'age': (age)}, (label)), steps=num_steps)
    self.assertEqual(initial_global_step + num_steps,
                     est.get_variable_value(mock_opt.iterations.name))
    self._assert_checkpoint(
        n_classes,
        expected_global_step=initial_global_step + num_steps,
        expected_age_weight=age_weight,
        expected_bias=bias)

def testBinaryClassesFromCheckpointMultiBatch(self):
    self._testFromCheckpointMultiBatch(n_classes=2)

def testMultiClassesFromCheckpointMultiBatch(self):
    self._testFromCheckpointMultiBatch(n_classes=4)

class BaseLinearClassifierEvaluationTest(object):

def __init__(self, linear_classifier_fn, fc_lib=feature_column_v2):
    self._linear_classifier_fn = linear_classifier_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

def _test_evaluation_for_simple_data(self, n_classes):
    label = 1
    age = 1.

# For binary case, the expected weight has shape (1,1). For multi class
    # case, the shape is (1, n_classes). In order to test the weights, set
    # weights as 2.0 * range(n_classes).
    age_weight = [[-11.0]] if n_classes == 2 else (np.reshape(
        -11.0 * np.array(list(range(n_classes)), dtype=np.float32),
        (1, n_classes)))
    bias = [-30.0] if n_classes == 2 else [-30.0] * n_classes

with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

est = self._linear_classifier_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        model_dir=self._model_dir)
    eval_metrics = est.evaluate(
        input_fn=lambda: ({
            'age': ((age,),)
        }, ((label,),)), steps=1)

if n_classes == 2:
      # Binary classes: loss = sum(corss_entropy(41)) = 41.
      expected_metrics = {
          metric_keys.MetricKeys.LOSS: 41.,
          tf.compat.v1.GraphKeys.GLOBAL_STEP: 100,
          metric_keys.MetricKeys.LOSS_MEAN: 41.,
          metric_keys.MetricKeys.ACCURACY: 0.,
          metric_keys.MetricKeys.PRECISION: 0.,
          metric_keys.MetricKeys.RECALL: 0.,
          metric_keys.MetricKeys.PREDICTION_MEAN: 0.,
          metric_keys.MetricKeys.LABEL_MEAN: 1.,
          metric_keys.MetricKeys.ACCURACY_BASELINE: 1,
          metric_keys.MetricKeys.AUC: 0.,
          metric_keys.MetricKeys.AUC_PR: 1.,
      }
    else:
      # Multi classes: loss = 1 * -log ( soft_max(logits)[label] )
      logits = age_weight * age + bias
      logits_exp = np.exp(logits)
      softmax = logits_exp / logits_exp.sum()
      expected_loss = -1 * math.log(softmax[0, label])

expected_metrics = {
          metric_keys.MetricKeys.LOSS: expected_loss,
          metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
          tf.compat.v1.GraphKeys.GLOBAL_STEP: 100,
          metric_keys.MetricKeys.ACCURACY: 0.,
      }

self.assertAllClose(
        sorted_key_dict(expected_metrics),
        sorted_key_dict(eval_metrics),
        rtol=1e-3)

def test_binary_classes_evaluation_for_simple_data(self):
    self._test_evaluation_for_simple_data(n_classes=2)

def test_multi_classes_evaluation_for_simple_data(self):
    self._test_evaluation_for_simple_data(n_classes=4)

def _test_evaluation_batch(self, n_classes):
    """Tests evaluation for batch_size==2."""
    label = [1, 0]
    age = [17., 18.]
    # For binary case, the expected weight has shape (1,1). For multi class
    # case, the shape is (1, n_classes). In order to test the weights, set
    # weights as 2.0 * range(n_classes).
    age_weight = [[2.0]] if n_classes == 2 else (np.reshape(
        2.0 * np.array(list(range(n_classes)), dtype=np.float32),
        (1, n_classes)))
    bias = [-35.0] if n_classes == 2 else [-35.0] * n_classes
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

if n_classes == 2:
      # Logits are (-1., 1.) labels are (1, 0).
      # Loss is
      #   loss for row 1: 1 * -log(sigmoid(-1)) = 1.3133
      #   loss for row 2: (1 - 0) * -log(1 - sigmoid(1)) = 1.3133
      expected_loss = (1.3133 * 2) / 2  # Divided by batch size

expected_metrics = {
          metric_keys.MetricKeys.LOSS: expected_loss,
          tf.compat.v1.GraphKeys.GLOBAL_STEP: 100,
          metric_keys.MetricKeys.LOSS_MEAN: expected_loss,
          metric_keys.MetricKeys.ACCURACY: 0.,
          metric_keys.MetricKeys.PRECISION: 0.,
          metric_keys.MetricKeys.RECALL: 0.,
          metric_keys.MetricKeys.PREDICTION_MEAN: 0.5,
          metric_keys.MetricKeys.LABEL_MEAN: 0.5,
          metric_keys.MetricKeys.ACCURACY_BASELINE: 0.5,
          metric_keys.MetricKeys.AUC: 0.,
          metric_keys.MetricKeys.AUC_PR: 0.3068,
      }
    else:
      # Multi classes: loss = 1 * -log ( soft_max(logits)[label] )
      logits = age_weight * np.reshape(age, (2, 1)) + bias
      logits_exp = np.exp(logits)
      softmax_row_0 = logits_exp[0] / logits_exp[0].sum()
      softmax_row_1 = logits_exp[1] / logits_exp[1].sum()
      expected_loss_0 = -1 * math.log(softmax_row_0[label[0]])
      expected_loss_1 = -1 * math.log(softmax_row_1[label[1]])
      expected_loss = (expected_loss_0 + expected_loss_1) / 2  # batch size

self.assertAllClose(
        sorted_key_dict(expected_metrics),
        sorted_key_dict(eval_metrics),
        rtol=1e-3)

def test_binary_classes_evaluation_batch(self):
    self._test_evaluation_batch(n_classes=2)

def test_multi_classes_evaluation_batch(self):
    self._test_evaluation_batch(n_classes=4)

def _test_evaluation_weights(self, n_classes):
    """Tests evaluation with weights."""

label = [1, 0]
    age = [17., 18.]
    weights = [1., 2.]
    # For binary case, the expected weight has shape (1,1). For multi class
    # case, the shape is (1, n_classes). In order to test the weights, set
    # weights as 2.0 * range(n_classes).
    age_weight = [[2.0]] if n_classes == 2 else (np.reshape(
        2.0 * np.array(list(range(n_classes)), dtype=np.float32),
        (1, n_classes)))
    bias = [-35.0] if n_classes == 2 else [-35.0] * n_classes
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

est = self._linear_classifier_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        n_classes=n_classes,
        weight_column='w',
        model_dir=self._model_dir)
    eval_metrics = est.evaluate(
        input_fn=lambda: ({
            'age': (age),
            'w': (weights)
        }, (label)), steps=1)

if n_classes == 2:
      # Logits are (-1., 1.) labels are (1, 0).
      # Loss is
      #   loss for row 1: 1 * -log(sigmoid(-1)) = 1.3133
      #   loss for row 2: (1 - 0) * -log(1 - sigmoid(1)) = 1.3133
      #   weights = [1., 2.]
      expected_loss = (1.3133 * (1. + 2.)) / 2  # Divided by batch size
      loss_mean = (1.3133 * (1. + 2.)) / (1.0 + 2.0)
      label_mean = np.average(label, weights=weights)
      logits = [-1, 1]
      logistics = sigmoid(np.array(logits))
      predictions_mean = np.average(logistics, weights=weights)

expected_metrics = {
          metric_keys.MetricKeys.LOSS: expected_loss,
          tf.compat.v1.GraphKeys.GLOBAL_STEP: 100,
          metric_keys.MetricKeys.LOSS_MEAN: loss_mean,
          metric_keys.MetricKeys.ACCURACY: 0.,
          metric_keys.MetricKeys.PRECISION: 0.,
          metric_keys.MetricKeys.RECALL: 0.,
          metric_keys.MetricKeys.PREDICTION_MEAN: predictions_mean,
          metric_keys.MetricKeys.LABEL_MEAN: label_mean,
          metric_keys.MetricKeys.ACCURACY_BASELINE:
              (max(label_mean, 1 - label_mean)),
          metric_keys.MetricKeys.AUC: 0.,
          metric_keys.MetricKeys.AUC_PR: 0.1891,
      }
    else:
      # Multi classes: unweighted_loss = 1 * -log ( soft_max(logits)[label] )
      logits = age_weight * np.reshape(age, (2, 1)) + bias
      logits_exp = np.exp(logits)
      softmax_row_0 = logits_exp[0] / logits_exp[0].sum()
      softmax_row_1 = logits_exp[1] / logits_exp[1].sum()
      expected_loss_0 = -1 * math.log(softmax_row_0[label[0]])
      expected_loss_1 = -1 * math.log(softmax_row_1[label[1]])
      loss_mean = np.average([expected_loss_0, expected_loss_1],
                             weights=weights)
      expected_loss = (loss_mean * np.sum(weights)) / 2  # batch size

expected_metrics = {
          metric_keys.MetricKeys.LOSS: expected_loss,
          metric_keys.MetricKeys.LOSS_MEAN: loss_mean,
          tf.compat.v1.GraphKeys.GLOBAL_STEP: 100,
          metric_keys.MetricKeys.ACCURACY: 0.,
      }

self.assertAllClose(
        sorted_key_dict(expected_metrics),
        sorted_key_dict(eval_metrics),
        rtol=1e-3)

def test_binary_classes_evaluation_weights(self):
    self._test_evaluation_weights(n_classes=2)

def test_multi_classes_evaluation_weights(self):
    self._test_evaluation_weights(n_classes=4)

class BaseLinearClassifierPredictTest(object):

def __init__(self, linear_classifier_fn, fc_lib=feature_column_v2):
    self._linear_classifier_fn = linear_classifier_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

def _testPredictions(self, n_classes, label_vocabulary, label_output_fn):
    """Tests predict when all variables are one-dimensional."""
    age = 1.

with tf.Graph().as_default():
      tf.Variable(age_weight, name=AGE_WEIGHT_NAME)
      tf.Variable(bias, name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

est = self._linear_classifier_fn(
        feature_columns=(self._fc_lib.numeric_column('age'),),
        label_vocabulary=label_vocabulary,
        n_classes=n_classes,
        model_dir=self._model_dir)

predict_input_fn = numpy_io.numpy_input_fn(
        x={'age': np.array([[age]])},
        y=None,
        batch_size=1,
        num_epochs=1,
        shuffle=False)
    predictions = list(est.predict(input_fn=predict_input_fn))

if n_classes == 2:
      scalar_logits = np.asscalar(
          np.reshape(np.array(age_weight) * age + bias, (1,)))
      two_classes_logits = [0, scalar_logits]
      two_classes_logits_exp = np.exp(two_classes_logits)
      softmax = two_classes_logits_exp / two_classes_logits_exp.sum()

expected_predictions = {
          'class_ids': [0],
          'all_class_ids': [0, 1],
          'classes': [label_output_fn(0)],
          'all_classes': [label_output_fn(0),
                          label_output_fn(1)],
          'logistic': [sigmoid(np.array(scalar_logits))],
          'logits': [scalar_logits],
          'probabilities': softmax,
      }
    else:
      onedim_logits = np.reshape(np.array(age_weight) * age + bias, (-1,))
      class_ids = onedim_logits.argmax()
      all_class_ids = list(range(len(onedim_logits)))
      logits_exp = np.exp(onedim_logits)
      softmax = logits_exp / logits_exp.sum()
      expected_predictions = {
          'class_ids': [class_ids],
          'all_class_ids': all_class_ids,
          'classes': [label_output_fn(class_ids)],
          'all_classes': [label_output_fn(i) for i in all_class_ids],
          'logits': onedim_logits,
          'probabilities': softmax,
      }

self.assertEqual(1, len(predictions))
    # assertAllClose cannot handle byte type.
    self.assertEqual(expected_predictions['classes'], predictions[0]['classes'])
    expected_predictions.pop('classes')
    predictions[0].pop('classes')
    self.assertAllEqual(expected_predictions['all_classes'],
                        predictions[0]['all_classes'])
    expected_predictions.pop('all_classes')
    predictions[0].pop('all_classes')
    self.assertAllClose(
        sorted_key_dict(expected_predictions), sorted_key_dict(predictions[0]))

def testBinaryClassesWithoutLabelVocabulary(self):
    n_classes = 2
    self._testPredictions(
        n_classes,
        label_vocabulary=None,
        label_output_fn=lambda x: ('%s' % x).encode())

def testBinaryClassesWithLabelVocabulary(self):
    n_classes = 2
    self._testPredictions(
        n_classes,
        label_vocabulary=['class_vocab_{}'.format(i) for i in range(n_classes)],
        label_output_fn=lambda x: ('class_vocab_%s' % x).encode())

def testMultiClassesWithoutLabelVocabulary(self):
    n_classes = 4
    self._testPredictions(
        n_classes,
        label_vocabulary=None,
        label_output_fn=lambda x: ('%s' % x).encode())

def testMultiClassesWithLabelVocabulary(self):
    n_classes = 4
    self._testPredictions(
        n_classes,
        label_vocabulary=['class_vocab_{}'.format(i) for i in range(n_classes)],
        label_output_fn=lambda x: ('class_vocab_%s' % x).encode())

feature_columns = (self._fc_lib.categorical_column_with_vocabulary_list(
        'language', vocabulary_list=['a', 'b', 'c']),)

# Check prediction for each sparse_combiner.
    # With sparse_combiner = 'sum', we have
    # logits_1 = w_a + w_c + bias
    #          = 2.0 + 5.0 + 5.0 = 12.0
    # logits_2 = w_b + w_c + bias
    #          = 3.0 + 5.0 + 5.0 = 13.0
    linear_classifier = self._linear_classifier_fn(
        feature_columns=feature_columns, model_dir=self._model_dir)
    predictions = linear_classifier.predict(input_fn=_input_fn)
    predicted_scores = list([x['logits'] for x in predictions])
    self.assertAllClose([[12.0], [13.0]], predicted_scores)

# With sparse_combiner = 'mean', we have
    # logits_1 = 1/2 * (w_a + w_c) + bias
    #          = 1/2 * (2.0 + 5.0) + 5.0 = 8.5
    # logits_2 = 1/2 * (w_b + w_c) + bias
    #          = 1/2 * (3.0 + 5.0) + 5.0 = 9.0
    linear_classifier = self._linear_classifier_fn(
        feature_columns=feature_columns,
        model_dir=self._model_dir,
        sparse_combiner='mean')
    predictions = linear_classifier.predict(input_fn=_input_fn)
    predicted_scores = list([x['logits'] for x in predictions])
    self.assertAllClose([[8.5], [9.0]], predicted_scores)

# With sparse_combiner = 'sqrtn', we have
    # logits_1 = sqrt(2)/2 * (w_a + w_c) + bias
    #          = sqrt(2)/2 * (2.0 + 5.0) + 5.0 = 9.94974
    # logits_2 = sqrt(2)/2 * (w_b + w_c) + bias
    #          = sqrt(2)/2 * (3.0 + 5.0) + 5.0 = 10.65685
    linear_classifier = self._linear_classifier_fn(
        feature_columns=feature_columns,
        model_dir=self._model_dir,
        sparse_combiner='sqrtn')
    predictions = linear_classifier.predict(input_fn=_input_fn)
    predicted_scores = list([x['logits'] for x in predictions])
    self.assertAllClose([[9.94974], [10.65685]], predicted_scores)

class BaseLinearClassifierIntegrationTest(object):

def __init__(self, linear_classifier_fn, fc_lib=feature_column_v2):
    self._linear_classifier_fn = linear_classifier_fn
    self._fc_lib = fc_lib

def setUp(self):
    self._model_dir = tempfile.mkdtemp()

def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

def _test_complete_flow(self, n_classes, train_input_fn, eval_input_fn,
                          predict_input_fn, input_dimension, prediction_length):
    feature_columns = [
        self._fc_lib.numeric_column('x', shape=(input_dimension,))
    ]
    est = self._linear_classifier_fn(
        feature_columns=feature_columns,
        n_classes=n_classes,
        model_dir=self._model_dir)

# TRAIN
    # learn y = x
    est.train(train_input_fn, steps=200)

# EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(200, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn(metric_keys.MetricKeys.LOSS, six.iterkeys(scores))

# PREDICT
    predictions = np.array(
        [x['classes'] for x in est.predict(predict_input_fn)])
    self.assertAllEqual((prediction_length, 1), predictions.shape)

def _test_numpy_input_fn(self, n_classes):
    """Tests complete flow with numpy_input_fn."""
    input_dimension = 4
    batch_size = 10
    prediction_length = batch_size
    data = np.linspace(0., 2., batch_size * input_dimension, dtype=np.float32)
    data = data.reshape(batch_size, input_dimension)
    target = np.array([1] * batch_size)

train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=target,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=target,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=None,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)

self._test_complete_flow(
        n_classes=n_classes,
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        prediction_length=prediction_length)

def test_binary_classes_numpy_input_fn(self):
    self._test_numpy_input_fn(n_classes=2)

def test_multi_classes_numpy_input_fn(self):
    self._test_numpy_input_fn(n_classes=4)

def _test_pandas_input_fn(self, n_classes):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return

# Pandas DataFrame natually supports 1 dim data only.
    input_dimension = 1
    batch_size = 10
    data = np.array([1., 2., 3., 4.], dtype=np.float32)
    target = np.array([1, 0, 1, 0], dtype=np.int32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(target)
    prediction_length = 4

def test_binary_classes_pandas_input_fn(self):
    self._test_pandas_input_fn(n_classes=2)

def test_multi_classes_pandas_input_fn(self):
    self._test_pandas_input_fn(n_classes=4)

def _test_input_fn_from_parse_example(self, n_classes):
    """Tests complete flow with input_fn constructed from parse_example."""
    input_dimension = 2
    batch_size = 10
    prediction_length = batch_size
    data = np.linspace(0., 2., batch_size * input_dimension, dtype=np.float32)
    data = data.reshape(batch_size, input_dimension)
    target = np.array([1] * batch_size, dtype=np.int64)

serialized_examples = []
    for x, y in zip(data, target):
      example = example_pb2.Example(
          features=feature_pb2.Features(
              feature={
                  'x':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(value=x)),
                  'y':
                      feature_pb2.Feature(
                          int64_list=feature_pb2.Int64List(value=[y])),
              }))
      serialized_examples.append(example.SerializeToString())

feature_spec = {
        'x': tf.io.FixedLenFeature([input_dimension], tf.dtypes.float32),
        'y': tf.io.FixedLenFeature([1], tf.dtypes.int64),
    }

self._test_complete_flow(
        n_classes=n_classes,
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=input_dimension,
        prediction_length=prediction_length)

def test_binary_classes_input_fn_from_parse_example(self):
    self._test_input_fn_from_parse_example(n_classes=2)

def test_multi_classes_input_fn_from_parse_example(self):
    self._test_input_fn_from_parse_example(n_classes=4)

class BaseLinearLogitFnTest(object):

def __init__(self, fc_lib=feature_column_v2):
    self._fc_lib = fc_lib

def test_basic_logit_correctness(self):
    """linear_logit_fn simply wraps feature_column_lib.linear_model."""
    age = self._fc_lib.numeric_column('age')
    with tf.Graph().as_default():
      logit_fn = linear.linear_logit_fn_builder(units=2, feature_columns=[age])
      logits = logit_fn(features={'age': [[23.], [31.]]})
      bias_var = tf.compat.v1.get_collection(
          tf.compat.v1.GraphKeys.GLOBAL_VARIABLES,
          'linear_model/bias_weights')[0]
      age_var = tf.compat.v1.get_collection(
          tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, 'linear_model/age')[0]
      with tf.compat.v1.Session() as sess:
        sess.run([tf.compat.v1.initializers.global_variables()])
        self.assertAllClose([[0., 0.], [0., 0.]], logits.eval())
        sess.run(bias_var.assign([10., 5.]))
        self.assertAllClose([[10., 5.], [10., 5.]], logits.eval())
        sess.run(age_var.assign([[2.0, 3.0]]))
        # [2 * 23 + 10, 3 * 23 + 5] = [56, 74].
        # [2 * 31 + 10, 3 * 31 + 5] = [72, 98]
        self.assertAllClose([[56., 74.], [72., 98.]], logits.eval())

def test_compute_fraction_of_zero_v2(self):
    """Tests the calculation of sparsity."""
    if self._fc_lib != feature_column_v2:
      return

age = tf.feature_column.numeric_column('age')
    occupation = tf.feature_column.categorical_column_with_hash_bucket(
        'occupation', hash_bucket_size=5)
    with tf.Graph().as_default():
      model = linear.LinearModel(
          feature_columns=[age, occupation], units=3, name='linear_model')
      features = {
          'age': [[23.], [31.]],
          'occupation': [['doctor'], ['engineer']]
      }
      model(features)
      variables = model.variables
      variables.remove(model.bias)
      fraction_zero = linear._compute_fraction_of_zero(variables)
      age_var = tf.compat.v1.get_collection(
          tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, 'linear_model/age')[0]
      with tf.compat.v1.Session() as sess:
        sess.run([tf.compat.v1.initializers.global_variables()])
        # Upon initialization, all variables will be zero.
        self.assertAllClose(1, fraction_zero.eval())

sess.run(age_var.assign([[2.0, 0.0, -1.0]]))
        # 1 of the 3 age weights are zero, and all of the 15 (5 hash buckets
        # x 3-dim output) are zero.
        self.assertAllClose(16. / 18., fraction_zero.eval())

class BaseLinearWarmStartingTest(object):

def __init__(self,
               _linear_classifier_fn,
               _linear_regressor_fn,
               fc_lib=feature_column_v2):
    self._linear_classifier_fn = _linear_classifier_fn
    self._linear_regressor_fn = _linear_regressor_fn
    self._fc_lib = fc_lib

def setUp(self):
    # Create a directory to save our old checkpoint and vocabularies to.
    self._ckpt_and_vocab_dir = tempfile.mkdtemp()

# Make a dummy input_fn.
    def _input_fn():
      features = {
          'age': [[23.], [31.]],
          'age_in_years': [[23.], [31.]],
          'occupation': [['doctor'], ['consultant']]
      }
      return features, [0, 1]

self._input_fn = _input_fn

def tearDown(self):
    # Clean up checkpoint / vocab dir.
    tf.compat.v1.summary.FileWriterCache.clear()
    shutil.rmtree(self._ckpt_and_vocab_dir)

def test_classifier_basic_warm_starting(self):
    """Tests correctness of LinearClassifier default warm-start."""
    age = self._fc_lib.numeric_column('age')

# Create a LinearClassifier and train to save a checkpoint.
    linear_classifier = self._linear_classifier_fn(
        feature_columns=[age],
        model_dir=self._ckpt_and_vocab_dir,
        n_classes=4,
        optimizer='SGD')
    linear_classifier.train(input_fn=self._input_fn, max_steps=1)

# Create a second LinearClassifier, warm-started from the first.  Use a
    # learning_rate = 0.0 optimizer to check values (use SGD so we don't have
    # accumulator values that change).
    warm_started_linear_classifier = self._linear_classifier_fn(
        feature_columns=[age],
        n_classes=4,
        optimizer=gradient_descent.SGD(learning_rate=0.0),
        warm_start_from=linear_classifier.model_dir)

warm_started_linear_classifier.train(input_fn=self._input_fn, max_steps=1)
    for variable_name in warm_started_linear_classifier.get_variable_names():
      # Learning rate is also checkpointed in V2 optimizer. So we need to make
      # sure it uses the new value after warm started.
      if 'learning_rate' in variable_name:
        self.assertAllClose(
            0.0,
            warm_started_linear_classifier.get_variable_value(variable_name))
      else:
        self.assertAllClose(
            linear_classifier.get_variable_value(variable_name),
            warm_started_linear_classifier.get_variable_value(variable_name))

def test_regressor_basic_warm_starting(self):
    """Tests correctness of LinearRegressor default warm-start."""
    age = self._fc_lib.numeric_column('age')

# Create a LinearRegressor and train to save a checkpoint.
    linear_regressor = self._linear_regressor_fn(
        feature_columns=[age],
        model_dir=self._ckpt_and_vocab_dir,
        optimizer='SGD')
    linear_regressor.train(input_fn=self._input_fn, max_steps=1)

# Create a second LinearRegressor, warm-started from the first.  Use a
    # learning_rate = 0.0 optimizer to check values (use SGD so we don't have
    # accumulator values that change).
    warm_started_linear_regressor = self._linear_regressor_fn(
        feature_columns=[age],
        optimizer=gradient_descent.SGD(learning_rate=0.0),
        warm_start_from=linear_regressor.model_dir)

warm_started_linear_regressor.train(input_fn=self._input_fn, max_steps=1)
    for variable_name in warm_started_linear_regressor.get_variable_names():
      # Learning rate is also checkpointed in V2 optimizer. So we need to make
      # sure it uses the new value after warm started.
      if 'learning_rate' in variable_name:
        self.assertAllClose(
            0.0,
            warm_started_linear_regressor.get_variable_value(variable_name))
      else:
        self.assertAllClose(
            linear_regressor.get_variable_value(variable_name),
            warm_started_linear_regressor.get_variable_value(variable_name))

def test_warm_starting_selective_variables(self):
    """Tests selecting variables to warm-start."""
    age = self._fc_lib.numeric_column('age')

# Create a second LinearClassifier, warm-started from the first.  Use a
    # learning_rate = 0.0 optimizer to check values (use SGD so we don't have
    # accumulator values that change).
    warm_started_linear_classifier = self._linear_classifier_fn(
        feature_columns=[age],
        n_classes=4,
        optimizer=gradient_descent.SGD(learning_rate=0.0),
        # The provided regular expression will only warm-start the age variable
        # and not the bias.
        warm_start_from=estimator.WarmStartSettings(
            ckpt_to_initialize_from=linear_classifier.model_dir,
            vars_to_warm_start='.*(age).*'))

warm_started_linear_classifier.train(input_fn=self._input_fn, max_steps=1)
    self.assertAllClose(
        linear_classifier.get_variable_value(AGE_WEIGHT_NAME),
        warm_started_linear_classifier.get_variable_value(AGE_WEIGHT_NAME))
    # Bias should still be zero from initialization.
    self.assertAllClose(
        [0.0] * 4, warm_started_linear_classifier.get_variable_value(BIAS_NAME))

def test_warm_starting_with_vocab_remapping_and_partitioning(self):
    """Tests warm-starting with vocab remapping and partitioning."""
    vocab_list = ['doctor', 'lawyer', 'consultant']
    vocab_file = os.path.join(self._ckpt_and_vocab_dir, 'occupation_vocab')
    with open(vocab_file, 'w') as f:
      f.write('\n'.join(vocab_list))
    occupation = self._fc_lib.categorical_column_with_vocabulary_file(
        'occupation',
        vocabulary_file=vocab_file,
        vocabulary_size=len(vocab_list))

# Create a LinearClassifier and train to save a checkpoint.
    linear_classifier = self._linear_classifier_fn(
        feature_columns=[occupation],
        model_dir=self._ckpt_and_vocab_dir,
        n_classes=4,
        optimizer='SGD')
    linear_classifier.train(input_fn=self._input_fn, max_steps=1)

# Create a second LinearClassifier, warm-started from the first.  Use a
    # learning_rate = 0.0 optimizer to check values (use SGD so we don't have
    # accumulator values that change).  Use a new FeatureColumn with a
    # different vocabulary for occupation.
    new_vocab_list = ['doctor', 'consultant', 'engineer']
    new_vocab_file = os.path.join(self._ckpt_and_vocab_dir,
                                  'new_occupation_vocab')
    with open(new_vocab_file, 'w') as f:
      f.write('\n'.join(new_vocab_list))
    new_occupation = self._fc_lib.categorical_column_with_vocabulary_file(
        'occupation',
        vocabulary_file=new_vocab_file,
        vocabulary_size=len(new_vocab_list))
    # We can create our VocabInfo object from the new and old occupation
    # FeatureColumn's.
    occupation_vocab_info = estimator.VocabInfo(
        new_vocab=new_occupation.vocabulary_file,
        new_vocab_size=new_occupation.vocabulary_size,
        num_oov_buckets=new_occupation.num_oov_buckets,
        old_vocab=occupation.vocabulary_file,
        old_vocab_size=occupation.vocabulary_size,
        # Can't use constant_initializer with load_and_remap.  In practice,
        # use a truncated normal initializer.
        backup_initializer=tf.compat.v1.initializers.random_uniform(
            minval=0.39, maxval=0.39))
    warm_started_linear_classifier = self._linear_classifier_fn(
        feature_columns=[occupation],
        n_classes=4,
        optimizer=gradient_descent.SGD(learning_rate=0.0),
        warm_start_from=estimator.WarmStartSettings(
            ckpt_to_initialize_from=linear_classifier.model_dir,
            var_name_to_vocab_info={
                OCCUPATION_WEIGHT_NAME: occupation_vocab_info
            },
            # Explicitly providing None here will only warm-start variables
            # referenced in var_name_to_vocab_info (the bias will not be
            # warm-started).
            vars_to_warm_start=None))

warm_started_linear_classifier.train(input_fn=self._input_fn, max_steps=1)
    # 'doctor' was ID-0 and still ID-0.
    self.assertAllClose(
        linear_classifier.get_variable_value(OCCUPATION_WEIGHT_NAME)[0, :],
        warm_started_linear_classifier.get_variable_value(
            OCCUPATION_WEIGHT_NAME)[0, :])
    # 'consultant' was ID-2 and now ID-1.
    self.assertAllClose(
        linear_classifier.get_variable_value(OCCUPATION_WEIGHT_NAME)[2, :],
        warm_started_linear_classifier.get_variable_value(
            OCCUPATION_WEIGHT_NAME)[1, :])
    # 'engineer' is a new entry and should be initialized with the
    # backup_initializer in VocabInfo.
    self.assertAllClose([0.39] * 4,
                        warm_started_linear_classifier.get_variable_value(
                            OCCUPATION_WEIGHT_NAME)[2, :])
    # Bias should still be zero (from initialization logic).
    self.assertAllClose(
        [0.0] * 4, warm_started_linear_classifier.get_variable_value(BIAS_NAME))

def test_warm_starting_with_naming_change(self):
    """Tests warm-starting with a Tensor name remapping."""
    age_in_years = self._fc_lib.numeric_column('age_in_years')

# Create a LinearClassifier and train to save a checkpoint.
    linear_classifier = self._linear_classifier_fn(
        feature_columns=[age_in_years],
        model_dir=self._ckpt_and_vocab_dir,
        n_classes=4,
        optimizer='SGD')
    linear_classifier.train(input_fn=self._input_fn, max_steps=1)

# Create a second LinearClassifier, warm-started from the first.  Use a
    # learning_rate = 0.0 optimizer to check values (use SGD so we don't have
    # accumulator values that change).
    warm_started_linear_classifier = self._linear_classifier_fn(
        feature_columns=[self._fc_lib.numeric_column('age')],
        n_classes=4,
        optimizer=gradient_descent.SGD(learning_rate=0.0),
        # The 'age' variable correspond to the 'age_in_years' variable in the
        # previous model.
        warm_start_from=estimator.WarmStartSettings(
            ckpt_to_initialize_from=linear_classifier.model_dir,
            var_name_to_prev_var_name={
                AGE_WEIGHT_NAME: AGE_WEIGHT_NAME.replace('age', 'age_in_years')
            }))

warm_started_linear_classifier.train(input_fn=self._input_fn, max_steps=1)
    self.assertAllClose(
        linear_classifier.get_variable_value(
            AGE_WEIGHT_NAME.replace('age', 'age_in_years')),
        warm_started_linear_classifier.get_variable_value(AGE_WEIGHT_NAME))
    # The bias is also warm-started (with no name remapping).
    self.assertAllClose(
        linear_classifier.get_variable_value(BIAS_NAME),
        warm_started_linear_classifier.get_variable_value(BIAS_NAME))