Python tf.compat.v1.tpu.shard用法及代码示例

用于并行执行的分片computation。

用法

tf.compat.v1.tpu.shard(
    computation:Callable[..., Any],
    inputs:Optional[List[core_types.Tensor]] = None,
    num_shards:int = 1,
    input_shard_axes:Optional[List[int]] = None,
    outputs_from_all_shards:Union[bool, List[bool]] = True,
    output_shard_axes:Optional[List[int]] = None,
    infeed_queue:Optional[tpu_feed.InfeedQueue] = None,
    device_assignment:Optional[tf.tpu.experimental.DeviceAssignment] = None,
    name:Optional[Text] = None,
    xla_options:Optional[tf.tpu.XLAOptions] = None
) -> List[core_types.Tensor]

参数

computation 一个 Python 函数，它构建一个计算以应用于输入的每个分片。
inputs 输入张量列表或无(相当于一个空列表)。每个输入张量都有一个对应的分片轴，由 input_shard_axes 给出，它的大小必须能被 num_shards 整除。
num_shards 分片的数量。
input_shard_axes 要沿其分片的维度列表 inputs 或 None 。 None 表示“沿维度 0 对所有输入进行分片”。如果不是 None ，则每个输入必须有一个维度。
outputs_from_all_shards 布尔值或布尔值列表。对于每个输出，如果 True ，则所有分片的输出都沿相应的 output_shard_axes 条目连接。否则，每个输出都取自任意分片。如果参数是布尔值，则参数的值用于每个输出。
output_shard_axes 连接 computation 或 None 输出的维度列表。 None 表示“沿维度 0 连接所有输出”。如果不是 None ，则每个输出必须有一个维度。如果 outputs_from_all_shards 为 False，则忽略。
infeed_queue 如果不是 None ，则使用 InfeedQueue 来增加 computation 的输入。
device_assignment 如果不是 None ，则使用 DeviceAssignment 说明计算中的逻辑核心与 TPU 拓扑中的物理核心之间的映射。如果 None ，则使用默认设备分配。 DeviceAssignment 如果计算的每个分片只使用一个核，或者只有一个分片，或者分片数等于 TPU 系统中的核数，则可以省略DeviceAssignment。
name (已弃用)什么都不做。
xla_options tpu.XLAOptions 的实例，指示传递给 XLA 编译器的选项。使用 None 作为默认选项。

输出张量列表。

抛出

ValueError 如果num_shards <= 0
ValueError 如果 len(input_shard_axes) != len(inputs)
ValueError 如果 len(output_shard_axes) != len(computation 的输出)

inputs 必须是张量列表或无(相当于一个空列表)，每个都有对应的分割轴(来自 input_shard_axes )。每个输入沿相应的轴拆分为num_shards 片段，计算并行应用于每个分片。

如果张量被 computation 词汇捕获，则张量将广播到所有分片。例如：，

x = tf.constant(7) def computation(): return x + 3 ... = shard(computation, ...)

作为输入。

如果 outputs_from_all_shards 为真，则 computation 的所有分片的输出将沿着它们的 output_shard_axes 连接在一起。否则，每个输出都取自任意分片。

计算的输入和输出必须至少是 rank-1 张量。

相关用法

注：本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.compat.v1.tpu.shard。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。

用法

参数

返回

抛出