總覽
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)
注意:如需有效控制代碼的詳細資訊,請參閱此處。