从 TF Hub 缓存模型下载

概览

目前,tensorflow_hub 库支持两种下载模型的模式。默认情况下,模型下载为压缩存档并缓存在磁盘上。其次,模型可以直接从远程存储中读入 TensorFlow。无论哪种方式,实际 Python 代码中对 tensorflow_hub 函数的调用都可以且应该继续使用模型的规范 tfhub.dev URL,这些 URL 在系统之间是可移植的,并且可用于浏览文档。在用户代码需要实际文件系统位置(下载和解压缩后,或将模型句柄解析为文件系统路径后)的罕见情况下,可以通过函数 hub.resolve(handle) 获取该位置。

压缩下载的缓存

当从 tfhub.dev(或其他 托管网站)下载并解压缩模型时,tensorflow_hub 库默认情况下会将模型缓存在文件系统上。建议在大多数环境中使用此模式,除非磁盘空间不足,但网络带宽和延迟极佳。

下载位置默认为本地临时目录,但可以通过设置环境变量 TFHUB_CACHE_DIR(推荐)或传递命令行标志 --tfhub_cache_dir 进行自定义。默认缓存位置 /tmp/tfhub_modules(或 os.path.join(tempfile.gettempdir(), "tfhub_modules") 的评估结果)在大多数情况下都应该有效。

希望在系统重新启动后进行持久缓存的用户可以将 TFHUB_CACHE_DIR 设置为其主目录中的位置。例如,Linux 系统上 bash shell 的用户可以在 ~/.bashrc 中添加如下行

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

...重新启动 shell,然后将使用此位置。使用持久位置时,请注意没有自动清理。

从远程存储中读取

用户可以指示 tensorflow_hub 库直接从远程存储(GCS)读取模型,而不是使用以下命令本地下载模型

os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"

或将命令行标志 --tfhub_model_load_format 设置为 UNCOMPRESSED。这样,无需缓存目录,这在磁盘空间不足但互联网连接速度快的情况下特别有用。

在 Colab 笔记本中使用 TPU 运行

colab.research.google.com 上,下载压缩模型会与 TPU 运行时冲突,因为计算工作负载被委托给另一台机器,而该机器默认情况下无法访问缓存位置。对于这种情况,有两种解决方法

1) 使用 TPU 工作器可以访问的 GCS 存储分区

最简单的解决方案是指示 tensorflow_hub 库从 TF Hub 的 GCS 存储分区读取模型,如上所述。拥有自己的 GCS 存储分区的用户可以改为指定其存储分区中的目录作为缓存位置,代码如下

import os
os.environ["TFHUB_CACHE_DIR"] = "gs://my-bucket/tfhub-modules-cache"

...在调用 tensorflow_hub 库之前。

2) 通过 Colab 主机重定向所有读取

另一种解决方法是通过 Colab 主机重定向所有读取(即使是大变量的读取)

load_options =
tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
reloaded_model = hub.load("https://tfhub.dev/...", options=load_options)

注意:有关有效句柄的更多信息,请参见 此处