Polish keras wrapped model init with feature columns (#2487)

* polish keras wrapped model init with feature columns

* update

* update

* update

Author: typhoonzero

python/sqlflow_submitter/tensorflow/diag.py

@@ -22,15 +22,16 @@ class SQLFlowDiagnostic(Exception):
     pass
 
 
-def check_and_load_estimator(estimator, model_params, warm_start_from=None):
+def load_pretrained_model_estimator(estimator,
+                                    model_params,
+                                    warm_start_from=None):
     if warm_start_from is not None:
         estimator_func = estimator.__init__ if inspect.isclass(
             estimator) else estimator
         estimator_spec = inspect.getargspec(estimator_func)
         # The constructor of Estimator contains named parameter "warm_start_from"
         warm_start_from_key = "warm_start_from"
         if warm_start_from_key in estimator_spec.args:
-            model_params = copy.copy(model_params)
             warm_start_from = os.path.abspath(warm_start_from)
 
             if is_tf_estimator(estimator):
@@ -48,6 +49,8 @@ def check_and_load_estimator(estimator, model_params, warm_start_from=None):
                 "Incremental training is not supported in {}".format(
                     estimator))
 
+
+def init_model(estimator, model_params):
     # load estimator class and diagnose the type error
     try:
         return estimator(**model_params)

python/sqlflow_submitter/tensorflow/evaluate.py

@@ -27,6 +27,7 @@
 
 from . import metrics
 from .input_fn import get_dataset_fn
+from .keras_with_feature_column_input import init_model_with_feature_column
 from .pai_distributed import define_tf_flags
 from .set_log_level import set_log_level
 
@@ -85,11 +86,12 @@ def evaluate(datasource,
             model_params["model_dir"] = FLAGS.checkpointDir
         else:
             model_params["model_dir"] = save
+        # tf estimator always have feature_column argument
         estimator = estimator_cls(**model_params)
         result_metrics = estimator_evaluate(estimator, eval_dataset,
                                             validation_metrics)
     else:
-        keras_model = estimator_cls(**model_params)
+        keras_model = init_model_with_feature_column(estimator, model_params)
         keras_model_pkg = sys.modules[estimator_cls.__module__]
         result_metrics = keras_evaluate(keras_model, eval_dataset, save,
                                         keras_model_pkg, validation_metrics)

python/sqlflow_submitter/tensorflow/explain.py

@@ -31,6 +31,7 @@
 
 from .get_tf_version import tf_is_version2
 from .input_fn import input_fn
+from .keras_with_feature_column_input import init_model_with_feature_column
 
 sns_colors = sns.color_palette('colorblind')
 # Disable Tensorflow INFO and WARNING logs
@@ -97,7 +98,9 @@ def _input_fn():
         return dataset.batch(1).cache()
 
     model_params.update(feature_columns)
-    estimator = estimator_cls(**model_params)
+
+    estimator = init_model_with_feature_column(estimator_cls, model_params)
+
     if estimator_cls in (tf.estimator.BoostedTreesClassifier,
                          tf.estimator.BoostedTreesRegressor):
         explain_boosted_trees(datasource, estimator, _input_fn, plot_type,

python/sqlflow_submitter/tensorflow/keras_with_feature_column_input.py

@@ -11,8 +11,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
+
 import tensorflow as tf
 
+from .diag import init_model
+
 
 class WrappedKerasModel(tf.keras.Model):
     def __init__(self, keras_model, model_params, feature_columns):
@@ -23,3 +27,23 @@ def __init__(self, keras_model, model_params, feature_columns):
     def __call__(self, inputs, training=True):
         x = self.feature_layer(inputs)
         return self.sub_model.__call__(x, training=training)
+
+
+def init_model_with_feature_column(estimator,
+                                   model_params,
+                                   has_none_optimizer=False):
+    """Check if estimator have argument "feature_column" and initialize the model
+       by wrapping the keras model if no "feature_column" argument detected.
+
+       NOTE: initalize estimator model can also use this function since estimators all have
+       "feature_column" argument.
+    """
+    argspec = inspect.getargspec(estimator)
+    if "feature_columns" not in argspec.args and not has_none_optimizer:
+        feature_columns = model_params["feature_columns"]
+        del model_params["feature_columns"]
+        classifier = WrappedKerasModel(estimator, model_params,
+                                       feature_columns)
+    else:
+        classifier = init_model(estimator, model_params)
+    return classifier

python/sqlflow_submitter/tensorflow/predict.py

@@ -32,7 +32,7 @@
 
 from .get_tf_version import tf_is_version2
 from .input_fn import get_dtype, parse_sparse_feature_predict, tf_generator
-from .keras_with_feature_column_input import WrappedKerasModel
+from .keras_with_feature_column_input import init_model_with_feature_column
 
 try:
     import sqlflow_models
@@ -57,19 +57,8 @@ def keras_predict(estimator, model_params, save, result_table, is_pai,
                   pai_table, feature_column_names, feature_metas,
                   result_col_name, datasource, select, hdfs_namenode_addr,
                   hive_location, hdfs_user, hdfs_pass):
-    signature = inspect.signature(estimator)
-    has_feature_columns_arg = False
-    for p in signature.parameters:
-        if signature.parameters[p].name == "feature_columns":
-            has_feature_columns_arg = True
-            break
-    if not has_feature_columns_arg:
-        feature_columns = model_params["feature_columns"]
-        del model_params["feature_columns"]
-        classifier = WrappedKerasModel(estimator, model_params,
-                                       feature_columns)
-    else:
-        classifier = estimator(**model_params)
+
+    classifier = init_model_with_feature_column(estimator, model_params)
     classifier_pkg = sys.modules[estimator.__module__]
     conn = None
     if is_pai:

python/sqlflow_submitter/tensorflow/train_estimator.py

@@ -17,7 +17,7 @@
 
 from ..model_metadata import save_model_metadata
 from . import metrics
-from .diag import check_and_load_estimator
+from .diag import init_model, load_pretrained_model_estimator
 from .get_tf_version import tf_is_version2
 from .input_fn import input_fn
 from .pai_distributed import make_estimator_distributed_runconfig
@@ -46,8 +46,9 @@ def estimator_train_and_save(estimator, model_params, save, is_pai, FLAGS,
         model_params["model_dir"] = save
 
     warm_start_from = save if load_pretrained_model else None
-    classifier = check_and_load_estimator(estimator, model_params,
-                                          warm_start_from)
+    if warm_start_from:
+        load_pretrained_model_estimator(estimator, model_params)
+    classifier = init_model(estimator, model_params)
 
     # do not add default Accuracy metric when using estimator to train, it will fail
     # when the estimator is a regressor, and estimator seems automatically add some

python/sqlflow_submitter/tensorflow/train_keras.py

@@ -23,10 +23,9 @@
 
 from ..model_metadata import save_model_metadata
 from . import metrics
-from .diag import check_and_load_estimator
 from .get_tf_version import tf_is_version2
 from .input_fn import input_fn
-from .keras_with_feature_column_input import WrappedKerasModel
+from .keras_with_feature_column_input import init_model_with_feature_column
 from .pai_distributed import (dump_into_tf_config,
                               make_distributed_info_without_evaluator)
 from .train_estimator import estimator_train_compiled
@@ -47,13 +46,6 @@ def keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
         loss = model_params["loss"]
         del model_params["loss"]
 
-    signature = inspect.signature(estimator)
-    has_feature_columns_arg = False
-    for p in signature.parameters:
-        if signature.parameters[p].name == "feature_columns":
-            has_feature_columns_arg = True
-            break
-
     classifier_pkg = sys.modules[estimator.__module__]
     # setting training metrics
     model_metrics = []
@@ -101,13 +93,8 @@ def keras_train_and_save(estimator, model_params, save, is_pai, FLAGS,
     else:
         validate_dataset = None
 
-    if not has_feature_columns_arg and not has_none_optimizer:
-        feature_columns = model_params["feature_columns"]
-        del model_params["feature_columns"]
-        classifier = WrappedKerasModel(estimator, model_params,
-                                       feature_columns)
-    else:
-        classifier = check_and_load_estimator(estimator, model_params)
+    classifier = init_model_with_feature_column(
+        estimator, model_params, has_none_optimizer=has_none_optimizer)
 
     # FIXME(sneaxiy): some models defined by other framework (not TensorFlow or XGBoost)
     # may return None optimizer.