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

构建运行复制的 TPU 计算的图运算符。

用法

tf.compat.v1.tpu.replicate(
    computation:Callable[..., Any],
    inputs:Optional[List[List[core_types.Tensor]]] = None,
    infeed_queue:Optional[tpu_feed.InfeedQueue] = None,
    device_assignment:Optional[tf.tpu.experimental.DeviceAssignment] = None,
    name:Optional[Text] = None,
    maximum_shapes:Optional[Any] = None,
    padding_spec:Optional[tf.compat.v1.tpu.PaddingSpec] = None,
    xla_options:Optional[tf.tpu.XLAOptions] = None
) -> List[Any]

参数

computation 一个 Python 函数，用于构建要复制的计算。
inputs 输入张量列表或 None(相当于 [[]] )的列表，由 [replica_num][input_num] 索引。所有副本必须具有相同数量的输入。每个输入都可以是一个嵌套结构，其中包含可转换为张量的值。请注意，传递兼容值的N-dimension 列表将产生标量张量的N-dimension 列表，而不是单个Rank-N 张量。如果您需要不同的行为，请使用 tf.convert_to_tensor 将部分输入转换为张量。
infeed_queue 如果不是 None ，则从 InfeedQueue 追加参数元组作为计算的输入。
device_assignment 如果不是 None ，则使用 DeviceAssignment 说明计算中的逻辑核心与 TPU 拓扑中的物理核心之间的映射。如果 None ，则使用默认设备分配。如果计算的每个副本只使用一个核，并且只有一个副本，或者副本数等于 TPU 系统中的核数，则可以省略DeviceAssignment。
name (已弃用)什么都不做。
maximum_shapes tf.TensorShape 的嵌套结构，表示每个副本中每个输入元素的相应组件应填充到的形状。任何未知维度(例如 tf.TensorShape 中的 tf.compat.v1.Dimension(None) 或 tensor-like 对象中的 -1)将在所有副本上填充到该维度的最大大小。 maximum_shapes 的结构需要与 inputs[0] 相同。
padding_spec tpu.PaddingSpec 指定的枚举。这说明了当inputs 到tpu.replicate 是动态的时的填充策略。一种用法是通过将值设置为 tpu.PaddingSpec.POWER_OF_TWO 来启用对输入的自动分桶，这有助于减少 XLA 端的重新编译。
xla_options tpu.XLAOptions 的实例，指示传递给 XLA 编译器的选项。使用 None 作为默认选项。

输出列表，索引为[replica_num]每个输出都可以是一个嵌套结构，与 computation() 返回的内容相同，但有一些例外。
例外情况包括：

1) 无输出：将返回 NoOp，计算时 control-depends。 2)单值输出：将返回一个包含该值的元组。 3) Operation-only 输出：将返回一个 NoOp，计算时会返回 control-depends。

抛出

ValueError 如果所有副本都没有相同数量的输入张量。
ValueError 如果每个副本的输入数量与 computation 的形式参数数量不匹配。
ValueError 如果静态 inputs 尺寸与 maximum_shapes 中给出的值不匹配。
ValueError 如果每个副本的输入结构与 maximum_shapes 的结构不匹配。

inputs 具有静态形状的基本用法示例：

def computation(x):
  x = x + 1
  return tf.math.reduce_mean(x)

x = tf.convert_to_tensor([1., 2., 3.])
y = tf.convert_to_tensor([4., 5., 6.])
tf.compat.v1.tpu.replicate(computation, inputs=[[x], [y]])

如果 inputs 具有动态形状，并且您希望自动对输入进行分桶以避免 XLA 重新编译。请参阅下面的高级示例：

def computation(x):
  x = x + 1
  return tf.math.reduce_mean(x)

# Assume input tensors in two replicas `x` and `y` both have dynamic shape
# ([None, 2]).
tf.compat.v1.tpu.replicate(
  computation,
  inputs=[x, y],
  maximum_shapes=[tf.TensorShape([None, None])],
  padding_spec=tf.compat.v1.tpu.PaddingSpec.POWER_OF_TWO)

相关用法

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

用法

参数

返回

抛出