tf.saved_model.save

Экспортирует объект Trackable obj в формат SavedModel.

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

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

class Adder(tf.Module):

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

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

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

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

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

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

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

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

Поскольку объекты tf.keras.Model также являются Trackable, эту функцию можно использовать для экспорта моделей Keras. Например, экспорт с указанной подписью:

class Model(tf.keras.Model):

  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
  def serve(self, serialized):
    ...

m = Model()
tf.saved_model.save(m, '/tmp/saved_model/')

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

class Model(tf.keras.Model):

  @tf.function
  def call(self, x):
    ...

m = Model()
tf.saved_model.save(
    m, '/tmp/saved_model/',
    signatures=m.call.get_concrete_function(
        tf.TensorSpec(shape=[None, 3], dtype=tf.float32, name="inp")))

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

x = input_layer.Input((4,), name="x")
y = core.Dense(5, name="out")(x)
model = training.Model(x, y)
tf.saved_model.save(model, '/tmp/saved_model/')
# The exported SavedModel takes "x" with shape [None, 4] and returns "out"
# with shape [None, 5]

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

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

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

Одна tf.function может генерировать множество ConcreteFunctions. Если инструмент на нижнем уровне хочет обратиться ко всем конкретным функциям, сгенерированным одной tf.function, вы можете использовать аргумент function_aliases для хранения отображения от псевдонима к всем именам конкретных функций. Например:

class MyModel:
@tf.function
def func():
  ...

@tf.function
def serve():
  ...
  func()

model = MyModel()
signatures = {
    'serving_default': model.serve.get_concrete_function(),
}
options = tf.saved_model.SaveOptions(function_aliases={
    'my_func': func,
})
tf.saved_model.save(model, export_dir, signatures, options)

Совместимость с Eager

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

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