`tf.model_to_estimator` raise AttributeError when I add a list object to keras subclass model

Question

I've created a keras subclass model just like this:

class SubModel(tf.Keras.Model):

    def __init__(self, features, **kwargs):
        """Init function of Model.
        Args:
            features: A list of SparseFeature and DenseFeature.
        """
        assert len(features) > 0
        super(SubModel, self).__init__(name='SubModel', **kwargs)
        self.features = features

Note that there is a features arrtibute in __init__ function which will be used in call method of this model. Everything works well when I train and evaluate the model with keras style.

But, now I want to convert this model to estimator using tf.keras.model_to_estimator function. It raise an error: AttributeError: '_ListWrapper' object has no attribute 'get_config'.

According to my debug, it's the features attribute which is added to the model cause this error. When convertint to estimator, it regard features as a layer of the model, and try to call the get_config function when cloning the model. It seems that all the attributes added to the model will be treated as layer when cloning the model.

But I really want to use features as a part of model, so that it can be accessed through other function of this model like call. Is there other ways to solve this?

Answer 1

I think tf.keras.model_to_estimator is compatible with Sequential or Functional API Keras model perfectly but poorly with Subclass model especially when implementing complicated operation in subclass.

So, If you have defined a subclass keras model, and want to covert it to estimator, the best way is define model_fn function, and put the keras model in it like code below:

def model_fn(features, labels, mode):
    model = SubModel()
    outputs = model(features)
    loss = tf.keras.losses.xx(labels, outputs)
    return tf.estimator.EstimatorSpec(...)