def test_schema_and_examples_are_save_correctly(saved_tf_iris_model): train_x, train_y = iris_data_utils.load_data()[0] X = pd.DataFrame(train_x) y = pd.Series(train_y) for signature in (None, infer_signature(X, y)): for example in (None, X.head(3)): with TempDir() as tmp: path = tmp.path("model") mlflow.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=path, signature=signature, input_example=example, ) mlflow_model = Model.load(path) assert signature == mlflow_model.signature if example is None: assert mlflow_model.saved_input_example_info is None else: assert all((_read_example(mlflow_model, path) == example).all())
def saved_tf_iris_model(tmpdir): # Following code from # https://github.com/tensorflow/models/blob/master/samples/core/get_started/premade_estimator.py train_x, train_y = iris_data_utils.load_data()[0] # Feature columns describe how to use the input. my_feature_columns = [] for key in train_x.keys(): my_feature_columns.append(tf.feature_column.numeric_column(key=key)) # Build 2 hidden layer DNN with 10, 10 units respectively. estimator = tf.estimator.DNNClassifier( feature_columns=my_feature_columns, # Two hidden layers of 10 nodes each. hidden_units=[10, 10], # The model must choose between 3 classes. n_classes=3, ) # Train the Model. batch_size = 100 train_steps = 1000 estimator.train( input_fn=lambda: iris_data_utils.train_input_fn( train_x, train_y, batch_size), steps=train_steps, ) # Generate predictions from the model predict_x = { "SepalLength": [5.1, 5.9, 6.9], "SepalWidth": [3.3, 3.0, 3.1], "PetalLength": [1.7, 4.2, 5.4], "PetalWidth": [0.5, 1.5, 2.1], } estimator_preds = estimator.predict( lambda: iris_data_utils.eval_input_fn(predict_x, None, batch_size)) # Building a dictionary of the predictions by the estimator. if sys.version_info < (3, 0): estimator_preds_dict = estimator_preds.next() else: estimator_preds_dict = next(estimator_preds) for row in estimator_preds: for key in row.keys(): estimator_preds_dict[key] = np.vstack( (estimator_preds_dict[key], row[key])) # Building a pandas DataFrame out of the prediction dictionary. estimator_preds_df = copy.deepcopy(estimator_preds_dict) for col in estimator_preds_df.keys(): if all(len(element) == 1 for element in estimator_preds_df[col]): estimator_preds_df[col] = estimator_preds_df[col].ravel() else: estimator_preds_df[col] = estimator_preds_df[col].tolist() # Building a DataFrame that contains the names of the flowers predicted. estimator_preds_df = pandas.DataFrame.from_dict(data=estimator_preds_df) estimator_preds_results = [ iris_data_utils.SPECIES[id[0]] for id in estimator_preds_dict["class_ids"] ] estimator_preds_results_df = pd.DataFrame( {"predictions": estimator_preds_results}) # Define a function for estimator inference feature_spec = {} for name in my_feature_columns: feature_spec[name.key] = tf.Variable([], dtype=tf.float64, name=name.key) receiver_fn = tf.estimator.export.build_raw_serving_input_receiver_fn( feature_spec) # Save the estimator and its inference function saved_estimator_path = str(tmpdir.mkdir("saved_model")) saved_estimator_path = estimator.export_saved_model( saved_estimator_path, receiver_fn).decode("utf-8") return SavedModelInfo( path=saved_estimator_path, meta_graph_tags=["serve"], signature_def_key="predict", inference_df=pd.DataFrame( data=predict_x, columns=[name.key for name in my_feature_columns]), expected_results_df=estimator_preds_results_df, raw_results=estimator_preds_dict, raw_df=estimator_preds_df, )