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)
  ...