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. Если найдено ровно один такой метод, он будет использован в качестве стандартной подписи для 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 на подкаталог SavedModel «variables/». В настоящее время переменные являются единственными поддерживаемыми tf.saved_model.save объектами со состоянием, но в будущем будут поддерживаться и другие (например, таблицы).

tf.function не жестко кодирует аннотации устройства извне тела функции, вместо этого используя контекст вызова устройства. Это означает, что, например, экспорт модели, работающей на GPU, и ее обслуживание на CPU, в целом, будет работать, с некоторыми исключениями. Аннотации 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 будет добавлять новые операции сохранения в стандартный граф на каждой итерации.

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