tf.saved_model.load_v2

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

tf.saved_model.load_v2(
    export_dir, tags=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, который позволяет извлекать функции для новых подграфов. Это эквивалентно импорту 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 для получения подробностей. Эти SavedModels также имеют атрибут .variables, содержащий импортированные переменные, и атрибут .graph, представляющий весь импортированный граф. Для SavedModels, экспортированных из tf.saved_model.save, переменные вместо этого назначаются атрибутам, которым они были назначены до экспорта.