Сохраненные модели из TF Hub в TensorFlow 2.

Формат SavedModel в TensorFlow 2 — это рекомендуемый способ поделиться предварительно обученными моделями и частями моделей в TensorFlow Hub. Он заменяет старый формат TF1 Hub и поставляется с новым набором API.

На этой странице объясняется, как повторно использовать сохраненные модели TF2 в программе TensorFlow 2 с низкоуровневым API hub.load() и его оболочкой hub.KerasLayer . (Обычно hub.KerasLayer объединяется с другими tf.keras.layers для построения модели Keras или model_fn оценщика TF2.) Эти API также могут загружать устаревшие модели в формате TF1 Hub, в определенных пределах, см. руководство по совместимости .

Пользователи TensorFlow 1 могут обновиться до TF 1.15, а затем использовать те же API. Старые версии TF1 не работают.

Использование SavedModels из TF Hub

Использование SavedModel в Керасе

Keras — это высокоуровневый API TensorFlow для построения моделей глубокого обучения путем составления объектов Keras Layer. Библиотека tensorflow_hub предоставляет hub.KerasLayer , который инициализируется URL-адресом (или путем к файловой системе) SavedModel, а затем обеспечивает вычисления из SavedModel, включая ее предварительно обученные веса.

Вот пример использования предварительно обученного встраивания текста:

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/nnlm-en-dim128/2"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

Исходя из этого, можно построить классификатор текста обычным способом Keras:

model = tf.keras.Sequential([
    embed,
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1, activation="sigmoid"),
])

Колаб по классификации текста — это полный пример обучения и оценки такого классификатора.

Веса модели в hub.KerasLayer по умолчанию установлены как необучаемые. См. раздел о тонкой настройке ниже, чтобы узнать, как это изменить. Веса распределяются между всеми приложениями одного и того же объекта слоя, как обычно в Keras.

Использование SavedModel в оценщике

Пользователи API-интерфейса Estimator TensorFlow для распределенного обучения могут использовать SavedModels из TF Hub, написав свой model_fn в hub.KerasLayer среди других tf.keras.layers .

За кулисами: загрузка и кеширование SavedModel

Использование SavedModel из TensorFlow Hub (или других серверов HTTPS, реализующих его протокол хостинга ) загружает и распаковывает ее в локальную файловую систему, если она еще не существует. Переменную среды TFHUB_CACHE_DIR можно установить для переопределения временного местоположения по умолчанию для кэширования загруженных и несжатых сохраненных моделей. Подробности см. в разделе Кэширование .

Использование SavedModel в низкоуровневом TensorFlow

Модельные ручки

SavedModels можно загрузить из указанного handle , где handle — это путь к файловой системе, действительный URL-адрес модели TFhub.dev (например, «https://tfhub.dev/...»). URL-адреса моделей Kaggle отражают обработку TFhub.dev в соответствии с нашими Условиями и лицензией, связанной с активами модели, например «https://www.kaggle.com/...». Дескрипторы из моделей Kaggle эквивалентны соответствующим дескрипторам TFhub.dev.

Функцияhub.load hub.load(handle) загружает и распаковывает SavedModel (если handle уже не является путем к файловой системе), а затем возвращает результат загрузки с помощью встроенной функции TensorFlow tf.saved_model.load() . Таким образом, hub.load() может обрабатывать любую действительную SavedModel (в отличие от своего предшественника hub.Module для TF1).

Расширенная тема: чего ожидать от SavedModel после загрузки

В зависимости от содержимого SavedModel результат obj = hub.load(...) может быть вызван различными способами (как более подробно описано в TensorFlow's SavedModel Guide :

  • Сигнатуры обслуживания SavedModel (если таковые имеются) представлены в виде словаря конкретных функций и могут вызываться как tensors_out = obj.signatures["serving_default"](**tensors_in) , со словарями тензоров, привязанными к соответствующим входным и выходным данным. имена и подчиняются ограничениям формы и типа подписи.

  • @tf.function -декорированные методы сохраненного объекта (если таковые имеются) восстанавливаются как объекты tf.function, которые могут быть вызваны всеми комбинациями тензорных и нетензорных аргументов, для которых tf.function была отслежена до сохранения. В частности, если существует метод obj.__call__ с подходящими трассировками, сам obj можно вызывать как функцию Python. Простой пример может выглядеть так: output_tensor = obj(input_tensor, training=False) .

Это оставляет огромную свободу в интерфейсах, которые могут реализовать SavedModels. Интерфейс Reusable SavedModels для obj устанавливает соглашения, согласно которым клиентский код, включая такие адаптеры, как hub.KerasLayer , знает, как использовать SavedModel.

Некоторые SavedModels могут не следовать этому соглашению, особенно целые модели, которые не предназначены для повторного использования в более крупных моделях, и просто предоставляют сигнатуры обслуживания.

Обучаемые переменные в SavedModel перезагружаются как обучаемые, и tf.GradientTape будет отслеживать их по умолчанию. См. раздел «Точная настройка» ниже, где приведены некоторые предостережения, и подумайте о том, чтобы для начала избежать этого. Даже если вы хотите провести тонкую настройку, вы можете посмотреть, советует ли obj.trainable_variables повторно обучать только подмножество изначально обучаемых переменных.

Создание сохраненных моделей для TF Hub

Обзор

SavedModel — это стандартный формат сериализации TensorFlow для обученных моделей или частей модели. Он хранит обученные веса модели вместе с точными операциями TensorFlow для выполнения вычислений. Его можно использовать независимо от кода, который его создал. В частности, его можно повторно использовать в различных API построения моделей высокого уровня, таких как Keras, поскольку операции TensorFlow являются их общим базовым языком.

Сохранение из Кераса

Начиная с TensorFlow 2, tf.keras.Model.save() и tf.keras.models.save_model() по умолчанию используют формат SavedModel (не HDF5). Полученные в результате SavedModels, которые можно использовать hub.load() , hub.KerasLayer и аналогичными адаптерами для других API высокого уровня по мере их появления.

Чтобы поделиться полной моделью Keras, просто сохраните ее с помощью include_optimizer=False .

Чтобы поделиться частью модели Keras, сделайте эту часть моделью, а затем сохраните ее. Вы можете либо выложить такой код с самого начала....

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(...)

…или вырежьте фрагмент, чтобы поделиться им постфактум (если он соответствует наслоению вашей полной модели):

full_model = tf.keras.Model(...)
sharing_input = full_model.get_layer(...).get_output_at(0)
sharing_output = full_model.get_layer(...).get_output_at(0)
piece_to_share = tf.keras.Model(sharing_input, sharing_output)
piece_to_share.save(..., include_optimizer=False)

Модели TensorFlow на GitHub используют первый подход для BERT (см. nlp/tools/export_tfhub_lib.py , обратите внимание на разделение между core_model для экспорта и pretrainer для восстановления контрольной точки) и последний подход для ResNet (см . Heritage/image_classification/tfhub_export.py ).

Сохранение из низкоуровневого TensorFlow

Для этого необходимо хорошее знание руководства SavedModel от TensorFlow.

Если вы хотите предоставить больше, чем просто подпись обслуживания, вам следует реализовать интерфейс Reusable SavedModel . Концептуально это выглядит так

class MyMulModel(tf.train.Checkpoint):
  def __init__(self, v_init):
    super().__init__()
    self.v = tf.Variable(v_init)
    self.variables = [self.v]
    self.trainable_variables = [self.v]
    self.regularization_losses = [
        tf.function(input_signature=[])(lambda: 0.001 * self.v**2),
    ]

  @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
  def __call__(self, inputs):
    return tf.multiply(inputs, self.v)

tf.saved_model.save(MyMulModel(2.0), "/tmp/my_mul")

layer = hub.KerasLayer("/tmp/my_mul")
print(layer([10., 20.]))  # [20., 40.]
layer.trainable = True
print(layer.trainable_weights)  # [2.]
print(layer.losses)  # 0.004

Тонкая настройка

Обучение уже обученных переменных импортированной SavedModel вместе с переменными модели вокруг нее называется точной настройкой SavedModel. Это может привести к повышению качества, но часто делает обучение более требовательным (может занять больше времени, больше зависеть от оптимизатора и его гиперпараметров, увеличить риск переобучения и потребовать увеличения набора данных, особенно для CNN). Мы советуем потребителям SavedModel заниматься тонкой настройкой только после установления хорошего режима обучения и только в том случае, если издатель SavedModel рекомендует это.

Точная настройка изменяет «непрерывные» параметры модели, которые обучаются. Он не меняет жестко запрограммированные преобразования, такие как токенизация ввода текста и сопоставление токенов с соответствующими записями в матрице внедрения.

Для потребителей SavedModel

Создание hub.KerasLayer вроде

layer = hub.KerasLayer(..., trainable=True)

обеспечивает точную настройку SavedModel, загруженной слоем. Он добавляет обучаемые веса и регуляризаторы весов, объявленные в SavedModel, к модели Keras и запускает вычисления SavedModel в режиме обучения (подумайте о выпадении и т. д.).

Совместная работа по классификации изображений содержит комплексный пример с дополнительной тонкой настройкой.

Реэкспорт результатов тонкой настройки

Опытные пользователи могут захотеть сохранить результаты тонкой настройки обратно в SavedModel, который можно будет использовать вместо первоначально загруженной модели. Это можно сделать с помощью кода типа

loaded_obj = hub.load("https://tfhub.dev/...")
hub_layer = hub.KerasLayer(loaded_obj, trainable=True, ...)

model = keras.Sequential([..., hub_layer, ...])
model.compile(...)
model.fit(...)

export_module_dir = os.path.join(os.getcwd(), "finetuned_model_export")
tf.saved_model.save(loaded_obj, export_module_dir)

Для создателей SavedModel

Создавая SavedModel для совместного использования в TensorFlow Hub, подумайте заранее, следует ли потребителям настраивать ее и как, и предоставьте рекомендации в документации.

Сохранение из модели Keras должно привести к тому, что все механизмы тонкой настройки (сохранение потерь при регуляризации веса, объявление обучаемых переменных, отслеживание __call__ как для training=True , так и training=False и т. д.)

Выберите интерфейс модели, который хорошо работает с градиентным потоком, например, выходные логиты вместо вероятностей softmax или прогнозов top-k.

Если в модели используется отсев, пакетная нормализация или аналогичные методы обучения, включающие гиперпараметры, установите для них значения, которые имеют смысл для многих ожидаемых целевых задач и размеров пакетов. (На момент написания этой статьи экономия на Keras не позволяет потребителям легко их корректировать.)

Регуляризаторы веса на отдельных слоях сохраняются (с их коэффициентами силы регуляризации), но регуляризация веса изнутри оптимизатора (например tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) теряется. Сообщите потребителям вашей SavedModel соответствующим образом.