概览
目前,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)
注意:有关有效句柄的更多信息,请参见 此处。