tf.train.Checkpoint

Управляет сохранением и восстановлением отслеживаемых значений на диске.

tf.train.Checkpoint(
    root=None, **kwargs
)

Используется в блокнотах

Объекты TensorFlow могут содержать отслеживаемое состояние, например tf.Variable, реализации tf.keras.optimizers.Optimizer, итераторы tf.data.Dataset, реализации tf.keras.Layer или реализации tf.keras.Model. Такие объекты называются отслеживаемыми объектами.

Объект Checkpoint можно создать для сохранения одного или группы отслеживаемых объектов в файл контрольной точки. Он поддерживает save_counter для нумерации контрольных точек.

Пример:

model = tf.keras.Model(...)
checkpoint = tf.train.Checkpoint(model)

# Save a checkpoint to /tmp/training_checkpoints-{save_counter}. Every time
# checkpoint.save is called, the save counter is increased.
save_path = checkpoint.save('/tmp/training_checkpoints')

# Restore the checkpointed values to the `model` object.
checkpoint.restore(save_path)

Пример 2:

import tensorflow as tf
import os

checkpoint_directory = "/tmp/training_checkpoints"
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")

# Create a Checkpoint that will manage two objects with trackable state,
# one we name "optimizer" and the other we name "model".
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
status = checkpoint.restore(tf.train.latest_checkpoint(checkpoint_directory))
for _ in range(num_training_steps):
  optimizer.minimize( ... )  # Variables will be restored on creation.
status.assert_consumed()  # Optional sanity checks.
checkpoint.save(file_prefix=checkpoint_prefix)

Checkpoint.save() и Checkpoint.restore() записывают и считывают контрольные точки на основе объектов, в отличие от tf.compat.v1.train.Saver в TensorFlow 1.x, который записывает и считывает контрольные точки на основе variable.name. При сохранении контрольных точек на основе объектов сохраняется граф зависимостей между объектами Python (Layer, Optimizer, Variable и т. д.) с именованными рёбрами. Этот граф используется для сопоставления переменных при восстановлении контрольной точки. Такой подход лучше переносит изменения в программе на Python и помогает поддерживать восстановление переменных при их создании.

Объекты Checkpoint зависят от объектов, переданных в конструкторы как именованные аргументы; каждой зависимости присваивается имя, совпадающее с именем соответствующего именованного аргумента. Классы TensorFlow, такие как Layer и Optimizer, автоматически добавляют зависимости от собственных переменных (например, «kernel» и «bias» для tf.keras.layers.Dense). Наследование от tf.keras.Model упрощает управление зависимостями в классах, определённых пользователем, поскольку Model перехватывает присваивание атрибутов. Например:

class Regress(tf.keras.Model):

  def __init__(self):
    super().__init__()
    self.input_transform = tf.keras.layers.Dense(10)
    # ...

  def call(self, inputs):
    x = self.input_transform(inputs)
    # ...

У этого Model есть зависимость с именем «input_transform» от его слоя Dense, который, в свою очередь, зависит от своих переменных. В результате при сохранении экземпляра Regress с помощью tf.train.Checkpoint также сохраняются все переменные, созданные слоем Dense.

Если переменные назначены нескольким рабочим узлам, каждый из них записывает собственный раздел контрольной точки. Затем эти разделы объединяются и переиндексируются так, чтобы образовать единую контрольную точку. Это позволяет не копировать все переменные на один рабочий узел, но требует, чтобы все рабочие узлы имели доступ к общей файловой системе.

Эта функция немного отличается от функции save_weights модели Keras. tf.keras.Model.save_weights создаёт файл контрольной точки с именем, указанным в filepath, а tf.train.Checkpoint нумерует контрольные точки, используя filepath в качестве префикса имён файлов контрольных точек. В остальном model.save_weights() и tf.train.Checkpoint(model).save() эквивалентны.

Подробности см. в руководстве по контрольным точкам обучения.

Методы

read

Посмотреть исходный код

read(
    save_path, options=None
)

Считывает контрольную точку обучения, записанную с помощью write.

Считывает этот объект Checkpoint и все объекты, от которых он зависит.

Этот метод похож на restore(), но не ожидает наличия в контрольной точке переменной save_counter. Он восстанавливает только те объекты, от которых контрольная точка уже зависит.

Метод предназначен главным образом для использования в высокоуровневых служебных средствах управления контрольными точками, которые используют write() вместо save() и применяют собственные механизмы нумерации и отслеживания контрольных точек.

Пример использования:

# Create a checkpoint with write()
ckpt = tf.train.Checkpoint(v=tf.Variable(1.))
path = ckpt.write('/tmp/my_checkpoint')

# Later, load the checkpoint with read()
# With restore() assert_consumed() would have failed.
checkpoint.read(path).assert_consumed()

# You can also pass options to read(). For example this
# runs the IO ops on the localhost:
options = tf.train.CheckpointOptions(
    experimental_io_device="/job:localhost")
checkpoint.read(path, options=options)

restore

Посмотреть исходный код

restore(
    save_path, options=None
)

Восстанавливает контрольную точку обучения.

Восстанавливает этот объект Checkpoint и все объекты, от которых он зависит.

Этот метод предназначен для загрузки контрольных точек, созданных с помощью save(). Для контрольных точек, созданных с помощью write(), используйте метод read(), который не ожидает переменную save_counter, добавляемую функцией save().

Функция restore() либо немедленно присваивает значения, если переменные для восстановления уже созданы, либо откладывает восстановление до момента создания переменных. Зависимости, добавленные после этого вызова, будут сопоставлены, если им соответствует объект в контрольной точке (запрос на восстановление будет ожидать в каждом отслеживаемом объекте, пока не будет добавлена ожидаемая зависимость).

checkpoint = tf.train.Checkpoint( ... )
checkpoint.restore(path)

# You can additionally pass options to restore():
options = tf.CheckpointOptions(experimental_io_device="/job:localhost")
checkpoint.restore(path, options=options)

Чтобы убедиться, что загрузка завершена и отложенных восстановлений больше не будет, используйте метод assert_consumed() объекта состояния, возвращённого функцией restore():

checkpoint.restore(path, options=options).assert_consumed()

Проверка вызовет ошибку, если в контрольной точке не найден какой-либо объект Python из графа зависимостей или если какое-либо сохранённое значение не соответствует объекту Python.

Контрольные точки TensorFlow 1.x на основе имён можно загружать с помощью этого метода через tf.compat.v1.train.Saver. Для сопоставления переменных используются имена. Как можно скорее преобразуйте контрольные точки на основе имён с помощью tf.train.Checkpoint.save.

Загрузка из контрольных точек SavedModel

Чтобы загрузить значения из SavedModel, просто передайте каталогу SavedModel в checkpoint.restore:

model = tf.keras.Model(...)
tf.saved_model.save(model, path)  # or model.save(path, save_format='tf')

checkpoint = tf.train.Checkpoint(model)
checkpoint.restore(path).expect_partial()

В этом примере для загруженного состояния вызывается expect_partial(), поскольку в контрольных точках SavedModel, сохранённых из Keras, часто создаются дополнительные ключи. В противном случае при завершении работы программа выводит множество предупреждений о неиспользуемых ключах.

save

Посмотреть исходный код

save(
    file_prefix, options=None
)

Сохраняет контрольную точку обучения и обеспечивает базовое управление контрольными точками.

Сохранённая контрольная точка содержит переменные, созданные этим объектом, а также все отслеживаемые объекты, от которых он зависит на момент вызова Checkpoint.save().

save — это простая удобная оболочка для метода write: она последовательно нумерует контрольные точки с помощью save_counter и обновляет метаданные, используемые функцией tf.train.latest_checkpoint. Другие служебные средства, также оборачивающие write и read, могут предоставлять более расширенные возможности управления контрольными точками, например сборку мусора и пользовательскую нумерацию. (Например, tf.train.CheckpointManager.)

step = tf.Variable(0, name="step")
checkpoint = tf.train.Checkpoint(step=step)
checkpoint.save("/tmp/ckpt")

# Later, read the checkpoint with restore()
checkpoint.restore("/tmp/ckpt-1")

# You can also pass options to save() and restore(). For example this
# runs the IO ops on the localhost:
options = tf.train.CheckpointOptions(experimental_io_device="/job:localhost")
checkpoint.save("/tmp/ckpt", options=options)

# Later, read the checkpoint with restore()
checkpoint.restore("/tmp/ckpt-1", options=options)

sync

Посмотреть исходный код

sync()

Ожидает завершения всех выполняющихся операций сохранения и восстановления.

write

Посмотреть исходный код

write(
    file_prefix, options=None
)

Записывает контрольную точку обучения.

Контрольная точка содержит переменные, созданные этим объектом, а также все отслеживаемые объекты, от которых он зависит на момент вызова Checkpoint.write().

write не нумерует контрольные точки, не увеличивает save_counter и не обновляет метаданные, используемые функцией tf.train.latest_checkpoint. Этот метод предназначен главным образом для высокоуровневых служебных средств управления контрольными точками. save предоставляет самую базовую реализацию этих возможностей.

Контрольные точки, записанные с помощью write, необходимо считывать с помощью read.

Пример использования:

step = tf.Variable(0, name="step")
checkpoint = tf.Checkpoint(step=step)
checkpoint.write("/tmp/ckpt")

# Later, read the checkpoint with read()
checkpoint.read("/tmp/ckpt")

# You can also pass options to write() and read(). For example this
# runs the IO ops on the localhost:
options = tf.CheckpointOptions(experimental_io_device="/job:localhost")
checkpoint.write("/tmp/ckpt", options=options)

# Later, read the checkpoint with read()
checkpoint.read("/tmp/ckpt", options=options)