Search code examples
pythontensorflowgoogle-cloud-storagetputensorflow-hub

loading tfdev.hub models in Tensorflow TPU from GCS

I am loading tfhub.dev model from GCS in colab TPU instance using

os.environ["TFHUB_CACHE_DIR"] = "gs://BUCKETNAME/model-cache-dir/"

with strategy.scope():
     layer = hub.KerasLayer("https://tfhub.dev/google/inaturalist/inception_v3/feature_vector/4",trainable=True)

but it takes a lot of time nearly 15 mins and also i get warning in the end

WARNING:absl:Deleting lock file gs://BUCKETNAME/model-cache-dir/a6cc63f37ce9d4a026a90b8d56f20a387de46a3f.lock due to inactivity.

any idea why is that

My guess : there is some sort of lock by tensorflow so that only one session can edit or modify the cache file , but after finishing the its operation hub.KerasLayer is not deleting the lock that causes inactivity.

Solution

  • This might be due to latency of copying files from GCS to GCS through the machine where colab is running.

    There is a way of using default /tmp location for TFHUB_CACHE_DIR that might be faster. Try not explicitly setting TFHUB_CACHED_DIR, and instead pass LoadOptions to hub.KerasLayer with experimental_io_device='/job:localhost', e.g.

    load_options = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost') layer = hub.KerasLayer(..., load_options=load_options)