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。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。

用法

參數

返回

拋出