Using a Keras model inside a TF estimator

Question

I want to use one of the pre-built keras' models (vgg, inception, resnet, etc) included in tf.keras.application for feature extraction to save me some time training.

What is the correct way to do this inside of an estimator model function?

This is what I currently have.

import tensorflow as tf

def model_fn(features, labels, mode):

    # Import the pretrained model
    base_model = tf.keras.applications.InceptionV3(
            weights='imagenet', 
            include_top=False,
            input_shape=(200,200,3)
    )

    # get the output features from InceptionV3
    resnet_features = base_model.predict(features['x'])

    # flatten and feed into dense layers
    pool2_flat = tf.layers.flatten(resnet_features)

    dense1 = tf.layers.dense(inputs=pool2_flat, units=5120, activation=tf.nn.relu)

    # ... Add in N number of dense layers depending on my application

    logits = tf.layers.dense(inputs=denseN, units=5)

    # Calculate Loss
    onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=5)

    loss = tf.losses.softmax_cross_entropy(
    onehot_labels=onehot_labels, logits=logits)

    optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    train_op = optimizer.minimize(
        loss=loss,
        global_step=tf.train.get_global_step()
    )

    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

if __name__ == "__main__":

    # import Xtrain and Ytrain

    classifier = tf.estimator.Estimator(
        model_fn=model_fn, model_dir="/tmp/conv_model")

    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={'x': Xtrain},
        y=Ytrain,
        batch_size=100,
        num_epochs=None,
        shuffle=True)

    classifier.train(
        input_fn=train_input_fn,
        steps=100)

However, this code throws the error:

TypeError: unsupported operand type(s) for /: 'Dimension' and 'float'

at line resnet_features = base_model.predict(features['x'])

I think this is because the keras model is expecting a numpy array, but the estimator is passing in a tf.Tensor.

So, what is the correct way to use a keras model inside of an estimator. And, if you're not suppose to do this, what is the simplest way to leverage a pre-trained model for transfer learning in TF?

Answer 1

I am not aware of any available method allowing you to create custom model_fn from pretrained keras model. An easier way is to use tf.keras.estimator.model_to_estimator()

model = tf.keras.applications.ResNet50(
    input_shape=(224, 224, 3),
    include_top=False,
    pooling='avg',
    weights='imagenet')
logits =  tf.keras.layers.Dense(10, 'softmax')(model.layers[-1].output)
model = tf.keras.models.Model(model.inputs, logits)
model.compile('adam', 'categorical_crossentropy', ['accuracy'])

# Convert Keras Model to tf.Estimator
estimator = tf.keras.estimator.model_to_estimator(keras_model=model)
estimator.train(input_fn=....)

However, if you would like to create custom model_fn to add more ops (e.g. Summary ops), you can write as following:

import tensorflow as tf

_INIT_WEIGHT = True

def model_fn(features, labels, mode, params):
  global _INIT_WEIGHT

  # This is important, it allows keras model to update weights
  tf.keras.backend.set_learning_phase(mode == tf.estimator.ModeKeys.TRAIN)

  model = tf.keras.applications.MobileNet(
      input_tensor=features,
      include_top=False,
      pooling='avg',
      weights='imagenet' if _INIT_WEIGHT else None)

  # Only init weights on first run
  if _INIT_WEIGHT:
    _INIT_WEIGHT = False

  feature_map = model(features)
  logits = tf.keras.layers.Dense(units=params['num_classes'])(feature_map)

  # loss
  loss = tf.losses.softmax_cross_entropy(labels=labels, logits=logits)
  ...

Answer 2

You can have only tensors in model_fn. Maybe you can try something like this. This can be considered as a hack. The better part is that this code apart from just providing model_fn, it also stores weights of the loaded model as a checkpoint in . This helps you to get the weights when you call estimator.train(...) or estimator.evaluate(...) from the checkpoint.

def model_fn(features, labels, mode):  

    # Import the pretrained model
    base_model = tf.keras.applications.InceptionV3(
        weights='imagenet', 
        include_top=False,
        input_shape=(200,200,3)
    )    

    # some check
    if not hasattr(m, 'optimizer'):
        raise ValueError(
            'Given keras model has not been compiled yet. '
            'Please compile first '
            'before creating the estimator.')

    # get estimator object from model
    keras_estimator_obj = tf.keras.estimator.model_to_estimator(
        keras_model=base_model,
        model_dir=<model_dir>,
        config=<run_config>,
    ) 

    # pull model_fn that we need (hack)
    return keras_estimator_obj._model_fn