Кэширование загрузки модели из TF Hub

Обзор

Библиотека tensorflow_hub на данный момент поддерживает два режима загрузки моделей. По умолчанию модель загружается в виде сжатого архива и кэшируется на диске. Во-вторых, модели можно напрямую считывать из удаленного хранилища в TensorFlow. В любом случае вызовы функций tensorflow_hub в реальном коде Python могут и должны продолжать использовать канонические URL-адреса моделей tfhub.dev, которые переносимы между системами и доступны для навигации по документации. В том редком случае, когда пользовательскому коду требуется фактическое местоположение файловой системы (после загрузки и распаковки или после преобразования дескриптора модели в путь файловой системы), его можно получить с помощью hub.resolve(handle) .

Кэширование сжатых загрузок

Библиотека tensorflow_hub по умолчанию кэширует модели в файловой системе, когда они были загружены с tfhub.dev (или других хостингов ) и распакованы. Этот режим рекомендуется для большинства сред, за исключением случаев, когда дискового пространства недостаточно, но пропускная способность сети и задержка превосходны.

По умолчанию местом загрузки является локальный временный каталог, но его можно настроить, задав переменную среды TFHUB_CACHE_DIR (рекомендуется) или передав флаг командной строки --tfhub_cache_dir . Местоположение кэша по умолчанию /tmp/tfhub_modules (или то, что оценивается os.path.join(tempfile.gettempdir(), "tfhub_modules") ) должно работать в большинстве случаев.

Пользователи, которые предпочитают постоянное кэширование при перезагрузке системы, могут вместо этого установить TFHUB_CACHE_DIR в свой домашний каталог. Например, пользователь оболочки bash в системе Linux может добавить в ~/.bashrc следующую строку:

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

... перезапустите оболочку, и тогда это местоположение будет использовано. При использовании постоянного местоположения имейте в виду, что автоматическая очистка не производится.

Чтение из удаленного хранилища

Пользователи могут поручить библиотеке tensorflow_hub напрямую считывать модели из удаленного хранилища (GCS) вместо загрузки моделей локально с помощью

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

или установив для флага командной строки --tfhub_model_load_format значение UNCOMPRESSED . Таким образом, каталог для кэширования не требуется, что особенно полезно в средах с небольшим дисковым пространством, но с быстрым подключением к Интернету.

Работа на ТПУ в ноутбуках Colab

На colab.research.google.com загрузка сжатых моделей будет конфликтовать со средой выполнения TPU, поскольку рабочая нагрузка по вычислениям делегируется другому компьютеру, который по умолчанию не имеет доступа к местоположению кэша. Есть два обходных пути для этой ситуации:

1) Используйте корзину GCS, к которой может получить доступ работник TPU.

Самое простое решение — поручить библиотеке tensorflow_hub читать модели из корзины GCS TF Hub, как описано выше. Пользователи с собственным сегментом 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)

Примечание. Дополнительную информацию о допустимых дескрипторах см . здесь .