tf.saved_model.save

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

tf.saved_model.save(
    obj, export_dir, signatures=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 на подкаталог "variables/" SavedModel. В настоящее время tf.saved_model.save поддерживает только переменные как единственные состоятельные объекты, но в будущем будут поддерживаться и другие (например, таблицы).

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

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

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

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

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