tensorflow2.0 - 如何离线加载 hub.KerasLayer？

Question

我已经通过 tensorflow hub 构建了一个模型并保存了它。
但是当我加载它时，我必须添加custom_objects={'KerasLayer':hub.KerasLayer} .
它将连接网络。

model = tf.keras.models.load_model('my_model.h5',custom_objects={'KerasLayer':hub.KerasLayer})

如何下载 hub.KerasLayer并离线加载？

Answer 1

您还可以使用 TFHUB_CACHE_DIR 自动执行此操作

import tensorflow_hub as hub
import os

os.environ["TFHUB_CACHE_DIR"] = "/tmp/model"
hub.KerasLayer("https://tfhub.dev/google/nnlm-id-dim50-with-normalization/2")

模型 Assets 将下载到/tmp/model/
future 调用 hub.KerasLayer("https://tfhub.dev/google/nnlm-id-dim50-with-normalization/2")将使用本地副本