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

用法

參數

返回

拋出