tf.saved_model.save

Экспортирует tf.Module (и подклассы) obj в формат SavedModel.

tf.saved_model.save(
    obj,
    export_dir: str,
    signatures=None,
    options: tf.saved_model.SaveOptions = None
)

Используется в ноутбуках

Объект obj должен наследоваться от класса Trackable.

Пример использования:

class Adder(tf.Module):
  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')

Полученный SavedModel будет доступен с именем входа "x", скалярным значением с типом float32.

Подписи

Подписи определяют типы входных и выходных данных для вычисления. Необязательный аргумент save signatures управляет тем, какие методы в obj будут доступны программам, которые используют SavedModel, например, серверным API. Функции Python могут быть декорированы с помощью @tf.function(input_signature=...) и передаваться в качестве подписей напрямую или отложенно с вызовом get_concrete_function к методу, декорированному @tf.function.

Пример:

class Adder(tf.Module):
  @tf.function
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(
  model, '/tmp/adder',signatures=model.add.get_concrete_function(
    tf.TensorSpec([], tf.float32)))

Если у @tf.function нет входной подписи и get_concrete_function не вызывается для этого метода, функция не будет напрямую вызываемой в восстановленной модели SavedModel.

Пример:

class Adder(tf.Module):
  @tf.function
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')
restored = tf.saved_model.load('/tmp/adder')
restored.add(1.)
Traceback (most recent call last):

ValueError: Found zero restored functions for caller function.

Если аргумент signatures опущен, obj будет искать методы, декорированные @tf.function. Если найдено ровно одно прослеженное @tf.function, этот метод будет использоваться в качестве подписи по умолчанию для SavedModel. В противном случае, любые @tf.function, прикрепленные к obj или его зависимостям, будут экспортированы для использования с tf.saved_model.load.

При вызове подписи в экспортированной SavedModel, аргументы Tensor идентифицируются по имени. Эти имена по умолчанию будут взяты из имен аргументов Python-функции. Их можно переопределить, указав аргумент name=... в соответствующем объекте tf.TensorSpec. Явное именование необходимо, если несколько Tensor передаются через один аргумент в Python-функцию.

Выходные данные функций, используемых в качестве signatures, должны быть либо плоскими списками, в этом случае выходные данные будут пронумерованы, либо словарем, сопоставляющим строковые ключи с Tensor, в этом случае ключи будут использованы для именования выходных данных.

Подписи доступны в объектах, возвращаемых tf.saved_model.load, в качестве атрибута .signatures. Это зарезервированный атрибут: tf.saved_model.save на объекте с пользовательским атрибутом .signatures вызовет исключение.

_Использование tf.savedmodel.save с моделями Keras

Хотя Keras имеет собственный API для сохранения и загрузки, эта функция может использоваться для экспорта моделей Keras. Например, экспорт с указанной подписью:

class Adder(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
  def concat(self, x):
     return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')

Экспорт из функции без фиксированной подписи:

class Adder(tf.keras.Model):
  @tf.function
  def concat(self, x):
     return x + x

model = Adder()
tf.saved_model.save(
  model, '/tmp/adder',
  signatures=model.concat.get_concrete_function(
    tf.TensorSpec(shape=[], dtype=tf.string, name="string_input")))

tf.keras.Model экземпляры, построенные из входных и выходных данных, уже имеют подпись и поэтому не требуют декоратора @tf.function или аргумента signatures. Если ни то, ни другое не указано, экспортируется проход модели вперёд.

x = tf.keras.layers.Input((4,), name="x")
y = tf.keras.layers.Dense(5, name="out")(x)
model = tf.keras.Model(x, y)
tf.saved_model.save(model, '/tmp/saved_model/')

Экспортируемый SavedModel принимает "x" с формой [None, 4] и возвращает "out" с формой [None, 5]

Переменные и контрольные точки

Переменные должны отслеживаться, назначая их в атрибут отслеживаемого объекта или непосредственно в атрибут obj. Объекты TensorFlow (например, слои из tf.keras.layers, оптимизаторы из tf.train) автоматически отслеживают свои переменные. Это та же схема отслеживания, что и у tf.train.Checkpoint, и экспортированный объект Checkpoint может быть восстановлен как контрольная точка обучения, указав tf.train.Checkpoint.restore на поддиректорию SavedModel's "variables/".

tf.function не жестко кодирует аннотации устройства извне тела функции, вместо этого используя контекст вызова. Это означает, например, что экспорт модели, работающей на GPU, и ее использование на CPU, как правило, будет работать, за некоторыми исключениями:

• tf.device аннотации внутри тела функции будут жестко закодированы в экспортированной модели; этот тип аннотации не рекомендуется.
• Операции, специфичные для устройства, например, с "cuDNN" в имени или с специфичными для устройства макетами, могут вызвать проблемы.
• В случае ConcreteFunctions активные стратегии распределения приведут к жесткой кодировке размещения устройств в функции.

SavedModels, экспортированные с помощью tf.saved_model.save, автоматически удаляют атрибуты с заданными по умолчанию значениями, что устраняет один источник несовместимостей, когда потребитель SavedModel использует более старую версию TensorFlow, чем производитель. Однако есть и другие источники несовместимостей, которые не обрабатываются автоматически, такие как случаи, когда экспортированная модель содержит операции, для которых у потребителя нет определений.

Не поддерживается при построении графа. С TensorFlow 1.x, tf.compat.v1.enable_eager_execution() должен быть вызван первым. Вызов tf.saved_model.save в цикле при построении графа с TensorFlow 1.x будет добавлять новые операции сохранения в стандартный граф на каждой итерации.

Нельзя вызывать изнутри тела функции.