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 можно восстановить как контрольную точку обучения, указав путь к подкаталогу "variables/" в SavedModel в tf.train.Checkpoint.restore.

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 на каждой итерации в граф по умолчанию будут добавляться новые операции сохранения.

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