tf.saved_model.load

Загрузка SavedModel из export_dir.

tf.saved_model.load(
    export_dir, tags=None, options=None
)

Подписи, связанные с SavedModel, доступны в качестве функций:

imported = tf.saved_model.load(path)
f = imported.signatures["serving_default"]
print(f(x=tf.constant([[1.]])))

Объекты, экспортированные с помощью tf.saved_model.save, дополнительно имеют присвоенные атрибуты отслеживаемых объектов и функций:

exported = tf.train.Checkpoint(v=tf.Variable(3.))
exported.f = tf.function(
    lambda x: exported.v * x,
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
tf.saved_model.save(exported, path)
imported = tf.saved_model.load(path)
assert 3. == imported.v.numpy()
assert 6. == imported.f(x=tf.constant(2.)).numpy()

Загрузка моделей Keras

Модели Keras отслеживаются, поэтому их можно сохранить в SavedModel. Объект, возвращаемый tf.saved_model.load, не является объектом Keras (т. е. не имеет методов .fit, .predict, и т. д.). Некоторые атрибуты и функции всё ещё доступны: .variables, .trainable_variables и .__call__.

model = tf.keras.Model(...)
tf.saved_model.save(model, path)
imported = tf.saved_model.load(path)
outputs = imported(inputs)

Используйте tf.keras.models.load_model для восстановления модели Keras.

Импорт SavedModel из TensorFlow 1.x

SavedModel из tf.estimator.Estimator или 1.x API SavedModel имеют плоский граф вместо объектов tf.function. Эти SavedModel будут загружены со следующими атрибутами:

• .signatures: Словарь, сопоставляющий имена подписей функциям.

• .prune(feeds, fetches): Метод, позволяющий извлекать функции для новых подграфов. Это эквивалентно импорту SavedModel и присвоению имен для входных и выходных данных в сессии из TensorFlow 1.x.

  imported = tf.saved_model.load(path_to_v1_saved_model)
  pruned = imported.prune("x:0", "out:0")
  pruned(tf.ones([]))
  

  Подробности см. в tf.compat.v1.wrap_function.

• .variables: Список импортированных переменных.

• .graph: Весь импортированный граф.

• .restore(save_path): Функция, восстанавливающая переменные из контрольной точки, сохранённой из tf.compat.v1.Saver.

Асинхронное использование SavedModel

При асинхронном использовании SavedModel (производитель — отдельный процесс) каталог SavedModel может появиться до завершения записи всех файлов, и tf.saved_model.load завершится с ошибкой, если указан неполный SavedModel. Вместо проверки каталога проверьте наличие файла "saved_model_dir/saved_model.pb". Этот файл записывается атомарно в качестве последней операции tf.saved_model.save.