Search code examples
tensorflowkerastpu

How to use Tensorflow Hub models on cloud TPU?

I'm trying to use a model from tensorflow hub on Kaggle. Like so:

m = tf.keras.Sequential([
hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4", output_shape=[1280],
               trainable=False),  # Can be True, see below.
tf.keras.layers.Dense(num_classes, activation='softmax')
])

m.build([None, 224, 224, 3])  # Batch input shape.

It works well with GPU, but as soon as I switch to TPU with TF records I get the following error:

InvalidArgumentError: Unsuccessful TensorSliceReader constructor: Failed to get matching files on /tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables: Unimplemented: File system scheme '[local]' not implemented (file: '/tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables')

However the set up and tfrecords dataset are all correct as it works with a switching the pretrained model to a keras application of the same model (i.e. for example above using the mobilenet keras application).

I tried caching but I have been unsuccessful, is there something I have to beware when following this guide: https://www.tensorflow.org/hub/caching

Thanks in advance!

Solution

  • The failure happens because the TPU is trying to load the TFHub model from /tmp/ which it doesn't have access to. You should be able to get this to work by:

    with strategy.scope():
  load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
  m = tf.keras.Sequential([
    hub.KerasLayer(
      "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4",
      output_shape=[1280],
      load_options=load_locally,
      trainable=False),  # Can be True, see below.
    tf.keras.layers.Dense(num_classes, activation='softmax')
  ])

    Source: EfficientNetB7 on 100+ flowers.