Trying to use TF-Hub module with CoLab TPU

[ayulockin]

I am doing some experimentation with style transfer and I used this TF-Hub module.

The accompanied CoLab Notebook can help reproduce the issue.

hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'
hub_module = hub.load(hub_handle)

This works fine with CoLab GPU. But I need to use the same with CoLab TPU. With CoLab TPU I am getting this error.

---------------------------------------------------------------------------

InvalidArgumentError                      Traceback (most recent call last)

<ipython-input-5-1b40bb7e9d3a> in <module>()
      2 
      3 hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'
----> 4 hub_module = hub.load(hub_handle)

16 frames

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

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

What's causing this issue? What's a potential way out?

Thanks in advance. :D

[gowthamkpr]

@ayulockin I tried to reproduce this issue on Google Cloab using TPU but I am not running into any error. Can you share a github gist showing its failing because I am not able to reproduce it. Thanks!

[ayulockin]

@gowthamkpr here's the GitHub gist showing the error. Gist

Hope this helps reproduce the error on CoLab TPU.

[kempy]

As discussed by email..

Unfortunately this is a known issue. Colab TPUs do not have access to the local file system which TF-Hub relies on.

TF-Hub modules are stored compressed on GCS and then downloaded, uncompressed, and cached locally. Because the Colab TPUs cannot read the local filesystem, you cannot read the downloaded model. A current workaround for this is to cache the uncompressed model on a public GCS bucket. You need to cache the uncompressed model. The system variable TFHUB_CACHE_DIR controls where the model gets cached

So, you should be able to set the cache dir like this:

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

And then call hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2") normally

If that doesn't work for some reason, then this should be roughly equivalent:

!mkdir /tmp/module
!curl -L "https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2?tf-hub-format=compressed" | tar -zxvC /tmp/module
!gsutil cp -R /tmp/module  "gs://my-bucket/model-cache-dir"

gs://my-bucket/model-cache-dir should now contain the uncompressed model which you can read directly with hub.load

And then you can call: hub.load("gs://my-bucket/model-cache-dir")

We are looking at ways to make sure that users don't need this workaround in a future release.

[Lawrence-Krukrubo]

The problem still persists, can't download feature vectors from tf_hub via Colab on TPU...

[WGierke]

Hi @Lawrence-Krukrubo
We recently added the possibility to directly read uncompressed models from GCS. This allows the TPU worker to also access the model files, while this is not possible when downloading compressed models and caching them on the TPU server (default behavior). The easiest fix is:

import os
import tensorflow_hub as hub

os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"
hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'
hub_module = hub.load(hub_handle)

You can find other solutions here. Thanks!