Благодарим за настройку Google I/O. Посмотреть все сеансы по запросу Смотреть по запросу

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

Обзор

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

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

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

Местом загрузки по умолчанию является локальный временный каталог, но его можно настроить, установив переменную среды 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://hub.tensorflow.google.cn/...", options=load_options)

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