How to load TF hub model from local system

Viewed 10617

One way is to download the model each time from tensorflow_hub like following

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/tf2-preview/nnlm-en-dim128/1"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

I want to download the file once and use again and again with out downloading each time

4 Answers
  1. Download your model from url + "?tf-hub-format=compressed"
    e.g. "https://tfhub.dev/google/tf2-preview/nnlm-en-dim128/1?tf-hub-format=compressed"
  2. Untar
  3. Load untarred folder in code
import tensorflow as tf
import tensorflow_hub as hub

embed = hub.KerasLayer('path/to/untarred/folder')
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

You can use the hub.load() method to load a TF Hub module. Also, the docs say,

Currently this method is fully supported only with TensorFlow 2.x and with modules created by calling tensorflow.saved_model.save(). The method works in both eager and graph modes.

The hub.load method has an argument handle. The types of modules handles are,

  1. Smart URL resolvers such as tfhub.dev, e.g.: https://tfhub.dev/google/nnlm-en-dim128/1.

  2. A directory on a file system supported by Tensorflow containing module files. This may include a local directory (e.g. /usr/local/mymodule) or a Google Cloud Storage bucket (gs://mymodule).

  3. A URL pointing to a TGZ archive of a module, e.g. https://example.com/mymodule.tar.gz.

You can use the 2nd and the 3rd points.

if anyone is wondering where the model is saved by default on windows, like me, it's here.

C:\Users\AvrakDavra\AppData\Local\Temp\tfhub_modules\

Obviously you can download anywhere and mention that path and tfhub will take from there, but in case. To instantly open the temp forlder on windows.

  1. Press WindowsButton+R
  2. Write %TEMP%

It'll open the temp folder for your username, and there the tfhub_modules folder be by default. It'll contain folders as below

enter image description here

Content of the text file are similar to below.

Module: https://tfhub.dev/google/universal-sentence-encoder/4 Download Time: 2021-07-17 18:17:09.714147 Downloader Hostname: LAPTOP(PID:12720)

Perhaps others may benefit from a concrete, reproducible answer. This post corresponds to this specific tfhub model.

tensorflow_hub version: 0.12.0
tensorflow version: 2.2.0

I set up the following path on my Linux server:

# Note, I manually created this entire path before ever downloading tfhub models
/opt/tfhub/tf2/bert_en_uncased_L-12_H-768_A-12_4/

(For various reasons, we have some needs for Tensorflow 1.x still, so I figured it might be a good idea to separate models based on if they are designed to work with tensorflow 1.x vs tensorflow 2.x, hence the tf2 in my path)

I then downloaded the model file, pushed it to my Linux server, placed it in the above location, and exectued:

# bash
tar xzf bert_en_uncased_L-12_H-768_A-12_4.tar.gz

That gave me the following files:

# python
import os
os.listdir("/opt/tfhub/tf2/bert_en_uncased_L-12_H-768_A-12_4/")
>>> ['keras_metadata.pb', 'saved_model.pb', 'assets', 'variables']

So then I can load the model like so:

# python
import tensorflow_hub as tfhub
import tensorflow as tf
bert_layer = tfhub.KerasLayer(tfhub.load("/opt/tfhub/tf2/bert_en_uncased_L-12_H-768_A-12_4"))
Related