tf.train.experimental.ShardingCallback

Функция обратного вызова для разбиения контрольной точки на шарды, а также её текстовое описание.

Обёртка для функции обратного вызова, которая выполняется, чтобы определить, как разделить тензоры на шарды, когда средство сохранения записывает шарды контрольной точки на диск.

Функция обратного вызова принимает на вход список tf.train.experimental.ShardableTensors (а также любые аргументы именованных параметров, определённые подклассом tf.train.experimental.ShardingCallback), и распределяет входные тензоры по разным шардам. Сначала тензоры группируются по задаче устройства (см. tf.DeviceSpec), затем функция обратного вызова вызывается для каждой группы тензоров.

При создании пользовательской функции обратного вызова необходимо учитывать несколько ограничений:

• Тензоры нельзя удалять из контрольной точки.
• Форму тензоров нельзя изменять.
• Типы данных тензоров нельзя изменять.
• Тензоры в одном шарде должны относиться к одной и той же задаче. После выполнения функции обратного вызова будут проведены проверки, чтобы убедиться, что эти ограничения не нарушены.

Пример простой пользовательской функции обратного вызова:

# Place all tensors in a single shard.
class AllInOnePolicy(tf.train.experimental.ShardingCallback):
  @property
  def description(self):
    return "Place all tensors in a single shard."

  def __call__(self, shardable_tensors):
    tensors = {}
    for shardable_tensor in shardable_tensors:
      tensor = shardable_tensor.tensor_save_spec.tensor
      checkpoint_key = shardable_tensor.checkpoint_key
      slice_spec = shardable_tensor.slice_spec

      tensors.set_default(checkpoint_key, {})[slice_spec] = tensor
    return [tensors]

ckpt.save(
    "path",
    options=tf.train.CheckpointOptions(
        experimental_sharding_callback=AllInOnePolicy()))

Атрибут description используется для идентификации функции обратного вызова и упрощения отладки при сохранении и восстановлении.

Чтобы передавать именованные аргументы, достаточно определить конструктор и передать их:

class ParameterPolicy(tf.train.experimental.ShardingCallback):
  def __init__(self, custom_param):
    self.custom_param = custom_param
  ...

ckpt.save(
    "path",
    options=tf.train.CheckpointOptions(
        experimental_sharding_callback=ParameterPolicy(custom_param=...)))

Методы

__call__

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

@abc.abstractmethod
__call__(
    shardable_tensors: Sequence[tf.train.experimental.ShardableTensor]
) -> Sequence[TensorSliceDict]

Вызов self как функции.