How do I save a tensorflow estimator created with tf.keras.estimator.model_to_estimator?

Question

How do I save a tensorflow estimator created with tf.keras.estimator.model_to_estimator?

The below is where I am at currently.

The keras model:

    def create_model(self):
        model = tf.keras.models.Sequential([
            tf.keras.layers.Reshape((self.num_features,), input_shape=(self.num_features, 1)),
            ...
            tf.keras.layers.Dense(1, activation='sigmoid')
        ])
        model.compile(optimizer='adam',
                      loss="binary_crossentropy",
                      metrics=['accuracy'])
        return model

From which I create the estimator:

self.model = self.create_model()
self.estimator = tf.keras.estimator.model_to_estimator(keras_model=self.model)
self.serving_input_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
tf.feature_column.make_parse_example_spec(
    [tf.feature_column.numeric_column(self.model.input_names[0], shape=self.model.layers[0].input_shape[1:])]
))

Then I train the model:

return tf.estimator.train_and_evaluate(
            estimator=self.estimator,
            train_spec=self.train_spec,
            eval_spec=self.test_spec)

Finally, I want to save this model:

self.estimator.export_saved_model(path, self.serving_input_receiver_fn)

This throws an error like the below. I can see that the path up to test_output/model\\temp-1599746280\\variables\\variables_temp_135e48a636f3441d8d35e49ff7fa2e67/ is created.

2020-09-10 21:58:01.338876: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at save_restore_v2_ops.cc:109 : Not found: Failed to create a NewWriteableFile: test_output/model\temp-1599746280\variables\variables_temp_135e48a636f3441d8d35e49ff7fa2e67/part-00000-of-00002.data-00000-of-00001.tempstate13255611844751724719 : The system cannot find the path specified.
; No such process

Answer 1

I'll answer my own question in case someone else is struggling with this. Note that I haven't got this working on a local dev environment as tf.estimator seems to have an issue writing to paths on my local windows environment, but it does work in SageMaker.

The code that works on SageMaker is below.

def __init__(..., model_dir)
...
    self.model = self.create_model()
    self.estimator = tf.keras.estimator.model_to_estimator(
        keras_model=self.model, 
        model_dir=model_dir
    )
...

def serving_input_fn(self):
    feature_spec = {
        self.model.input_names[0]: tf.io.FixedLenFeature(
            dtype=tf.float32, 
            shape=[self.num_features]
        )
    }
    logging.debug("feature spec: %s", feature_spec)
    return tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)()

...

def save_model(self, path):
    self.estimator.export_saved_model(path, self.serving_input_fn)

Note When running locally on windows to test I need to set model_dir = None in tf.keras.estimator.model_to_estimator and I can't save the model ie don't call self.estimator.export_saved_model