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.

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

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

• .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.

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

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