Сохраненные модели из 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

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

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

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://hub.tensorflow.google.cn/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 в оценщике

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

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

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

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

Ручки модели

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

Функция 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) , со словарями тензоров, заданными соответствующими входными и выходными данными. имена и с учетом ограничений формы и dtype подписи.

  • Декорированные @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 (см . legacy/image_classification/tfhub_export.py ).

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

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

Если вы хотите предоставить больше, чем просто обслуживающую подпись, вам следует реализовать интерфейс 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://hub.tensorflow.google.cn/...")
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 соответственно.