tf.train.experimental.MaxShardSizePolicy

Policy that splits tensors into shards with a max shard size.

Inherits From: ShardingCallback

tf.train.experimental.MaxShardSizePolicy(
    max_shard_size: int
)

Shards may exceed the max shard size if they contain 1. a single scalar/string tensor that could not be sliced and exceeds the max shard size or 2. the checkpoint object graph, whose size cannot be calculated when saving.

description

Methods

__call__

View source

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

Callback to split tensors into shards with a max shard size.

Args
shardable_tensors A list of ShardableTensors.

Returns
List of shard dicts containing tensors. [ {checkpoint key: {slice_spec: tensor} } ]