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.

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

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