Open source the NSL APIs to build TF Estimators that enable graph as well as adversarial regularization.

NSL Estimator APIs and test cases currently support only TF 1.x (and not TF 2.0 yet).

PiperOrigin-RevId: 273627369

Author: DualityGap

neural_structured_learning/BUILD

@@ -27,6 +27,7 @@ py_library(
     name = "neural_structured_learning",
     srcs = ["__init__.py"],
     deps = [
+        "//neural_structured_learning/estimator",
         "//neural_structured_learning/keras",
         "//neural_structured_learning/lib",
         "//neural_structured_learning/tools",

neural_structured_learning/__init__.py

@@ -1,6 +1,7 @@
 """Subpackages of Neural Structured Learning."""
 
 from neural_structured_learning import configs
+from neural_structured_learning import estimator
 from neural_structured_learning import keras
 from neural_structured_learning import lib
 from neural_structured_learning import tools

neural_structured_learning/estimator/BUILD

@@ -0,0 +1,85 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+
+# Description:
+# Build rules for Estimator APIs in Neural Structured Learning.
+
+# Placeholder for internal Python version compatibility macro.
+
+package(
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+exports_files(["LICENSE"])
+
+py_library(
+    name = "estimator",
+    srcs = ["__init__.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":adversarial_regularization",
+        ":graph_regularization",
+    ],
+)
+
+py_library(
+    name = "adversarial_regularization",
+    srcs = ["adversarial_regularization.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//neural_structured_learning/configs",
+        "//neural_structured_learning/lib",
+        # package tensorflow
+    ],
+)
+
+py_test(
+    name = "adversarial_regularization_test",
+    srcs = ["adversarial_regularization_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":estimator",
+        "//neural_structured_learning/configs",
+        # package numpy
+        # package tensorflow
+        # package tensorflow framework_test_lib,
+    ],
+)
+
+py_library(
+    name = "graph_regularization",
+    srcs = ["graph_regularization.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//neural_structured_learning/configs",
+        "//neural_structured_learning/lib:distances",
+        "//neural_structured_learning/lib:utils",
+        # package tensorflow
+    ],
+)
+
+py_test(
+    name = "graph_regularization_test",
+    srcs = ["graph_regularization_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":estimator",
+        # package protobuf,
+        "//neural_structured_learning/configs",
+        # package numpy
+        # package tensorflow
+        # package tensorflow framework_test_lib,
+    ],
+)

neural_structured_learning/estimator/README.md

@@ -0,0 +1,5 @@
+# Estimator APIs for Neural Structured Learning.
+
+The implementations here are provided on a strict "as is" basis, without
+warranties or conditions of any kind. Also, these implementations may not be
+compatible with certain TensorFlow versions (such as 2.0 or above).

neural_structured_learning/estimator/__init__.py

@@ -0,0 +1,8 @@
+r"""Estimator APIs for Neural Structured Learning.
+
+The current NSL Estimator APIs may not be compatible with certain TensorFlow
+versions (such as 2.0 or above).
+"""
+
+from neural_structured_learning.estimator.adversarial_regularization import add_adversarial_regularization
+from neural_structured_learning.estimator.graph_regularization import add_graph_regularization

neural_structured_learning/estimator/adversarial_regularization.py

@@ -0,0 +1,114 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+
+"""A wrapper function to enable adversarial regularization to an Estimator."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import neural_structured_learning.configs as nsl_configs
+import neural_structured_learning.lib as nsl_lib
+
+import tensorflow as tf
+
+
+def add_adversarial_regularization(estimator,
+                                   optimizer_fn=None,
+                                   adv_config=None):
+  """Adds adversarial regularization to a `tf.estimator.Estimator`.
+
+  Args:
+    estimator: An object of type `tf.estimator.Estimator`.
+    optimizer_fn: A function that accepts no arguments and returns an instance
+      of `tf.train.Optimizer`.
+    adv_config: An instance of `nsl.configs.AdvRegConfig` that specifies various
+      hyperparameters for adversarial regularization.
+
+  Returns:
+    A modified `tf.estimator.Estimator` object with adversarial regularization
+    incorporated into its loss.
+  """
+
+  if not adv_config:
+    adv_config = nsl_configs.AdvRegConfig()
+
+  base_model_fn = estimator._model_fn  # pylint: disable=protected-access
+
+  def adv_model_fn(features, labels, mode, params=None, config=None):
+    """The adversarial-regularized model_fn.
+
+    Args:
+      features: This is the first item returned from the `input_fn` passed to
+        `train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
+        or `dict` of same.
+      labels: This is the second item returned from the `input_fn` passed to
+        `train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
+        or dict of same (for multi-head models). If mode is
+        `tf.estimator.ModeKeys.PREDICT`, `labels=None` will be passed. If the
+        `model_fn`'s signature does not accept `mode`, the `model_fn` must still
+        be able to handle `labels=None`.
+      mode: Optional. Specifies if this is training, evaluation, or prediction.
+        See `tf.estimator.ModeKeys`.
+      params: Optional `dict` of hyperparameters. Will receive what is passed to
+        Estimator in the `params` parameter. This allows users to configure
+        Estimators from hyper parameter tuning.
+      config: Optional `estimator.RunConfig` object. Will receive what is passed
+        to Estimator as its `config` parameter, or a default value. Allows
+        setting up things in the model_fn based on configuration such as
+        `num_ps_replicas`, or `model_dir`. Unused currently.
+
+    Returns:
+      A `tf.EstimatorSpec` whose loss incorporates graph-based regularization.
+    """
+
+    # Uses the same variable scope for calculating the original objective and
+    # adversarial regularization.
+    with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(),
+                                     reuse=tf.compat.v1.AUTO_REUSE,
+                                     auxiliary_name_scope=False):
+      # If no 'params' is passed, then it is possible for base_model_fn not to
+      # accept a 'params' argument. See documentation for tf.estimator.Estimator
+      # for additional context.
+      if params:
+        original_spec = base_model_fn(features, labels, mode, params, config)
+      else:
+        original_spec = base_model_fn(features, labels, mode, config)
+
+      # Adversarial regularization only happens in training.
+      if mode != tf.estimator.ModeKeys.TRAIN:
+        return original_spec
+
+      adv_neighbor, _ = nsl_lib.gen_adv_neighbor(features, original_spec.loss,
+                                                 adv_config.adv_neighbor_config)
+
+      # Runs the base model again to compute loss on adv_neighbor.
+      adv_spec = base_model_fn(adv_neighbor, labels, mode, config)
+
+      final_loss = original_spec.loss + adv_config.multiplier * adv_spec.loss
+
+      if not optimizer_fn:
+        # Default to the Adagrad optimizer, the same as canned DNNEstimator.
+        optimizer = tf.train.AdagradOptimizer(learning_rate=0.05)
+      else:
+        optimizer = optimizer_fn()
+
+      final_train_op = optimizer.minimize(
+          loss=final_loss, global_step=tf.compat.v1.train.get_global_step())
+
+    return original_spec._replace(loss=final_loss, train_op=final_train_op)
+
+  # Replaces the model_fn while keeps other fields/methods in the estimator.
+  estimator._model_fn = adv_model_fn  # pylint: disable=protected-access
+  return estimator

neural_structured_learning/estimator/adversarial_regularization_test.py

@@ -0,0 +1,120 @@
+# Copyright 2019 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Tests for nsl.estimator.adversarial_regularization."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+import tempfile
+
+import neural_structured_learning.configs as nsl_configs
+import neural_structured_learning.estimator as nsl_estimator
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+FEATURE_NAME = 'x'
+WEIGHT_VARIABLE = 'linear/linear_model/' + FEATURE_NAME + '/weights'
+BIAS_VARIABLE = 'linear/linear_model/bias_weights'
+
+
+def single_batch_input_fn(features, labels=None):
+  def input_fn():
+    inputs = features if labels is None else (features, labels)
+    dataset = tf.data.Dataset.from_tensor_slices(inputs)
+    return dataset.batch(len(features))
+  return input_fn
+
+
+class AdversarialRegularizationTest(tf.test.TestCase):
+
+  def setUp(self):
+    super(AdversarialRegularizationTest, self).setUp()
+    self.model_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    if self.model_dir:
+      shutil.rmtree(self.model_dir)
+    super(AdversarialRegularizationTest, self).tearDown()
+
+  def build_linear_regressor(self, weight, bias):
+    with tf.Graph().as_default():
+      tf.Variable(weight, name=WEIGHT_VARIABLE)
+      tf.Variable(bias, name=BIAS_VARIABLE)
+      tf.Variable(100, name=tf.GraphKeys.GLOBAL_STEP, dtype=tf.int64)
+
+      with tf.Session() as sess:
+        sess.run([tf.global_variables_initializer()])
+        tf.train.Saver().save(sess, os.path.join(self.model_dir, 'model.ckpt'))
+
+    fc = tf.feature_column.numeric_column(FEATURE_NAME,
+                                          shape=np.array(weight).shape)
+    return tf.estimator.LinearRegressor(
+        feature_columns=(fc,), model_dir=self.model_dir, optimizer='SGD')
+
+  @test_util.run_v1_only('Requires tf.GraphKeys')
+  def test_adversarial_wrapper_not_affecting_predictions(self):
+    # base model: y = x + 2
+    base_est = self.build_linear_regressor(weight=[[1.0]], bias=[2.0])
+    adv_est = nsl_estimator.add_adversarial_regularization(base_est)
+    input_fn = single_batch_input_fn({FEATURE_NAME: np.array([[1.0], [2.0]])})
+    predictions = adv_est.predict(input_fn=input_fn)
+    predicted_scores = [x['predictions'] for x in predictions]
+    self.assertAllClose([[3.0], [4.0]], predicted_scores)
+
+  @test_util.run_v1_only('Requires tf.GraphKeys')
+  def test_adversarial_wrapper_adds_regularization(self):
+    # base model: y = w*x+b = 4*x1 + 3*x2 + 2
+    weight = np.array([[4.0], [3.0]], dtype=np.float32)
+    bias = np.array([2.0], dtype=np.float32)
+    x0, y0 = np.array([[1.0, 1.0]]), np.array([8.0])
+    adv_step_size = 0.1
+    learning_rate = 0.01
+
+    base_est = self.build_linear_regressor(weight=weight, bias=bias)
+    adv_config = nsl_configs.make_adv_reg_config(
+        multiplier=1.0,  # equal weight on original and adv examples
+        adv_step_size=adv_step_size)
+    adv_est = nsl_estimator.add_adversarial_regularization(
+        base_est,
+        optimizer_fn=lambda: tf.train.GradientDescentOptimizer(learning_rate),
+        adv_config=adv_config)
+    input_fn = single_batch_input_fn({FEATURE_NAME: x0}, y0)
+    adv_est.train(input_fn=input_fn, steps=1)
+
+    # Computes the gradients on original and adversarial examples.
+    orig_pred = np.dot(x0, weight) + bias  # [9.0]
+    orig_grad_w = 2 * (orig_pred - y0) * x0.T  # [[2.0], [2.0]]
+    orig_grad_b = 2 * (orig_pred - y0).reshape((1,))  # [2.0]
+    grad_x = 2 * (orig_pred - y0) * weight.T  # [[8.0, 6.0]]
+    perturbation = adv_step_size * grad_x / np.linalg.norm(grad_x)
+    x_adv = x0 + perturbation  # [[1.08, 1.06]]
+    adv_pred = np.dot(x_adv, weight) + bias  # [9.5]
+    adv_grad_w = 2 * (adv_pred - y0) * x_adv.T  # [[3.24], [3.18]]
+    adv_grad_b = 2 * (adv_pred - y0).reshape((1,))  # [3.0]
+
+    new_bias = bias - learning_rate * (orig_grad_b + adv_grad_b)
+    new_weight = weight - learning_rate * (orig_grad_w + adv_grad_w)
+    self.assertAllClose(new_bias, adv_est.get_variable_value(BIAS_VARIABLE))
+    self.assertAllClose(new_weight, adv_est.get_variable_value(WEIGHT_VARIABLE))
+
+
+if __name__ == '__main__':
+  tf.test.main()