Merge commit for internal changes

docs/tutorials/shape_constraints.ipynb

@@ -154,9 +154,9 @@
       },
       "outputs": [],
       "source": [
-        "NUM_EPOCHS = 500\n",
+        "NUM_EPOCHS = 1000\n",
         "BATCH_SIZE = 64\n",
-        "LEARNING_RATE=0.001"
+        "LEARNING_RATE=0.01"
       ]
     },
     {
@@ -376,9 +376,9 @@
         "\n",
         "# Generate datasets.\n",
         "np.random.seed(42)\n",
-        "data_train = sample_dataset(2000, testing_set=False)\n",
-        "data_val = sample_dataset(1000, testing_set=False)\n",
-        "data_test = sample_dataset(1000, testing_set=True)\n",
+        "data_train = sample_dataset(500, testing_set=False)\n",
+        "data_val = sample_dataset(500, testing_set=False)\n",
+        "data_test = sample_dataset(500, testing_set=True)\n",
         "\n",
         "# Plotting dataset densities.\n",
         "figsize(12, 5)\n",
@@ -536,9 +536,9 @@
         "    feature_columns=feature_columns,\n",
         "    # Hyper-params optimized on validation set.\n",
         "    n_batches_per_layer=1,\n",
-        "    max_depth=3,\n",
-        "    n_trees=20,\n",
-        "    min_node_weight=0.1,\n",
+        "    max_depth=2,\n",
+        "    n_trees=50,\n",
+        "    learning_rate=0.05,\n",
         "    config=tf.estimator.RunConfig(tf_random_seed=42),\n",
         ")\n",
         "gbt_estimator.train(input_fn=train_input_fn)\n",

examples/custom_estimators_uci_heart.py

@@ -18,7 +18,7 @@
 This example trains a TFL custom estimators on the UCI heart dataset.
 
 Example usage:
-custom_estimators_uci_heart --num_epochs=40
+custom_estimators_uci_heart --num_epochs=5000
 """
 
 from __future__ import absolute_import
@@ -38,7 +38,7 @@
 FLAGS = flags.FLAGS
 flags.DEFINE_float('learning_rate', 0.01, 'Learning rate.')
 flags.DEFINE_integer('batch_size', 100, 'Batch size.')
-flags.DEFINE_integer('num_epochs', 200, 'Number of training epoch.')
+flags.DEFINE_integer('num_epochs', 2000, 'Number of training epoch.')
 
 
 def main(_):
@@ -121,7 +121,7 @@ def model_fn(features, labels, mode, config):
         tfl.layers.PWLCalibration(
             input_keypoints=[0.0, 1.0, 2.0, 3.0],
             output_min=0.0,
-            output_max=lattice_sizes[0] - 1.0,
+            output_max=lattice_sizes[2] - 1.0,
             # You can specify TFL regularizers as tuple
             # ('regularizer name', l1, l2).
             kernel_regularizer=('hessian', 0.0, 1e-4),
@@ -130,16 +130,23 @@ def model_fn(features, labels, mode, config):
         tfl.layers.CategoricalCalibration(
             num_buckets=3,
             output_min=0.0,
-            output_max=lattice_sizes[1] - 1.0,
+            output_max=lattice_sizes[3] - 1.0,
             # Categorical monotonicity can be partial order.
             # (i, j) indicates that we must have output(i) <= output(i).
             # Make sure to set the lattice monotonicity to 1 for this dimension.
             monotonicities=[(0, 1), (0, 2)],
         )(inputs['thal']),
     ])
     output = tfl.layers.Lattice(
-        lattice_sizes=lattice_sizes, monotonicities=lattice_monotonicities)(
-            lattice_input)
+        lattice_sizes=lattice_sizes,
+        monotonicities=lattice_monotonicities,
+        # Add a kernel_initializer so that the Lattice is not initialized as a
+        # flat plane. The output_min and output_max could be arbitrary, as long
+        # as output_min < output_max.
+        kernel_initializer=tfl.lattice_layer.RandomMonotonicInitializer(
+            lattice_sizes=lattice_sizes, output_min=-10, output_max=10),
+    )(
+        lattice_input)
 
     training = (mode == tf.estimator.ModeKeys.TRAIN)
     model = tf.keras.Model(inputs=inputs, outputs=output)

setup.py

@@ -27,7 +27,7 @@
 # This version number should always be that of the *next* (unreleased) version.
 # Immediately after uploading a package to PyPI, you should increment the
 # version number and push to gitHub.
-__version__ = "2.0.7"
+__version__ = "2.0.8"
 
 if "--release" in sys.argv:
   sys.argv.remove("--release")

tensorflow_lattice/python/BUILD

@@ -117,6 +117,7 @@ py_library(
     deps = [
         ":categorical_calibration_layer",
         ":configs",
+        ":kronecker_factored_lattice_layer",
         ":lattice_layer",
         ":linear_layer",
         ":model_info",
@@ -197,6 +198,7 @@ py_test(
     srcs_version = "PY2AND3",
     deps = [
         ":kronecker_factored_lattice_layer",
+        ":kronecker_factored_lattice_lib",
         ":test_utils",
         # absl/logging dep,
         # absl/testing:parameterized dep,
@@ -329,6 +331,7 @@ py_library(
         ":aggregation_layer",
         ":categorical_calibration_layer",
         ":configs",
+        ":kronecker_factored_lattice_layer",
         ":lattice_layer",
         ":parallel_combination_layer",
         ":premade_lib",
@@ -346,6 +349,8 @@ py_library(
         ":aggregation_layer",
         ":categorical_calibration_layer",
         ":configs",
+        ":kronecker_factored_lattice_layer",
+        ":kronecker_factored_lattice_lib",
         ":lattice_layer",
         ":lattice_lib",
         ":linear_layer",
@@ -372,6 +377,7 @@ py_test(
         ":premade",
         ":premade_lib",
         # absl/logging dep,
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
     ],
@@ -439,6 +445,7 @@ py_library(
     srcs = ["rtl_layer.py"],
     srcs_version = "PY2AND3",
     deps = [
+        ":kronecker_factored_lattice_layer",
         ":lattice_layer",
         ":rtl_lib",
         # tensorflow:tensorflow_no_contrib dep,

tensorflow_lattice/python/configs.py

@@ -267,6 +267,8 @@ def __init__(self,
                num_lattices=None,
                lattice_rank=None,
                interpolation='hypercube',
+               parameterization='all_vertices',
+               num_terms=2,
                separate_calibrators=True,
                use_linear_combination=False,
                use_bias=False,
@@ -305,6 +307,34 @@ def __init__(self,
         'simplex' uses d+1 parameters and thus scales better. For details see
         `tfl.lattice_lib.evaluate_with_simplex_interpolation` and
         `tfl.lattice_lib.evaluate_with_hypercube_interpolation`.
+      parameterization: The parameterization of the lattice function class to
+        use. A lattice function is uniquely determined by specifying its value
+        on every lattice vertex. A parameterization scheme is a mapping from a
+        vector of parameters to a multidimensional array of lattice vertex
+        values. It can be one of:
+          - String `'all_vertices'`: This is the "traditional" parameterization
+            that keeps one scalar parameter per lattice vertex where the mapping
+            is essentially the identity map. With this scheme, the number of
+            parameters scales exponentially with the number of inputs to the
+            lattice. The underlying lattices used will be `tfl.layers.Lattice`
+            layers.
+          - String `'kronecker_factored'`: With this parameterization, for each
+            lattice input i we keep a collection of `num_terms` vectors each
+            having `feature_configs[0].lattice_size` entries (note that all
+            features must have the same lattice size). To obtain the tensor of
+            lattice vertex values, for `t=1,2,...,num_terms` we compute the
+            outer product of the `t'th` vector in each collection, multiply by a
+            per-term scale, and sum the resulting tensors. Finally, we add a
+            single shared bias parameter to each entry in the sum. With this
+            scheme, the number of parameters grows linearly with `lattice_rank`
+            (assuming lattice sizes and `num_terms` are held constant).
+            Currently, only monotonicity shape constraint and bound constraint
+            are supported for this scheme. Regularization is not currently
+            supported. The underlying lattices used will be
+            `tfl.layers.KroneckerFactoredLattice` layers.
+      num_terms: The number of terms in a lattice using `'kronecker_factored'`
+        parameterization. Ignored if parameterization is set to
+        `'all_vertices'`.
       separate_calibrators: If features should be separately calibrated for each
         lattice in the ensemble.
       use_linear_combination: If set to true, a linear combination layer will be
@@ -375,12 +405,15 @@ class CalibratedLatticeConfig(_Config, _HasFeatureConfigs,
   def __init__(self,
                feature_configs=None,
                interpolation='hypercube',
+               parameterization='all_vertices',
+               num_terms=2,
                regularizer_configs=None,
                output_min=None,
                output_max=None,
                output_calibration=False,
                output_calibration_num_keypoints=10,
-               output_initialization='quantiles'):
+               output_initialization='quantiles',
+               random_seed=0):
     """Initializes a `CalibratedLatticeConfig` instance.
 
     Args:
@@ -392,6 +425,34 @@ def __init__(self,
         'simplex' uses d+1 parameters and thus scales better. For details see
         `tfl.lattice_lib.evaluate_with_simplex_interpolation` and
         `tfl.lattice_lib.evaluate_with_hypercube_interpolation`.
+      parameterization: The parameterization of the lattice function class to
+        use. A lattice function is uniquely determined by specifying its value
+        on every lattice vertex. A parameterization scheme is a mapping from a
+        vector of parameters to a multidimensional array of lattice vertex
+        values. It can be one of:
+          - String `'all_vertices'`: This is the "traditional" parameterization
+            that keeps one scalar parameter per lattice vertex where the mapping
+            is essentially the identity map. With this scheme, the number of
+            parameters scales exponentially with the number of inputs to the
+            lattice. The underlying lattice used will be a `tfl.layers.Lattice`
+            layer.
+          - String `'kronecker_factored'`: With this parameterization, for each
+            lattice input i we keep a collection of `num_terms` vectors each
+            having `feature_configs[0].lattice_size` entries (note that all
+            features must have the same lattice size). To obtain the tensor of
+            lattice vertex values, for `t=1,2,...,num_terms` we compute the
+            outer product of the `t'th` vector in each collection, multiply by a
+            per-term scale, and sum the resulting tensors. Finally, we add a
+            single shared bias parameter to each entry in the sum. With this
+            scheme, the number of parameters grows linearly with
+            `len(feature_configs)` (assuming lattice sizes and `num_terms` are
+            held constant). Currently, only monotonicity shape constraint and
+            bound constraint are supported for this scheme. Regularization is
+            not currently supported. The underlying lattice used will be a
+            `tfl.layers.KroneckerFactoredLattice` layer.
+      num_terms: The number of terms in a lattice using `'kronecker_factored'`
+        parameterization. Ignored if parameterization is set to
+        `'all_vertices'`.
       regularizer_configs: A list of `tfl.configs.RegularizerConfig` instances
         that apply global regularization.
       output_min: Lower bound constraint on the output of the model.
@@ -410,6 +471,9 @@ def __init__(self,
           - String `'uniform'`: Output is initliazed uniformly in label range.
           - A list of numbers: To be used for initialization of the output
             lattice or output calibrator.
+      random_seed: Random seed to use for initialization of a lattice with
+        `'kronecker_factored'` parameterization. Ignored if parameterization is
+        set to `'all_vertices'`.
     """
     super(CalibratedLatticeConfig, self).__init__(locals())