資料類型¶

TorchRec 包含用於表示嵌入 (也稱為稀疏特徵) 的資料類型。稀疏特徵通常是指數，旨在饋入嵌入表。對於給定的批次，嵌入查找指數的數量是可變的。因此，需要一個參差不齊的維度來表示批次中可變數量的嵌入查找指數。

本節介紹用於表示稀疏特徵的三種 TorchRec 資料類型的類別：JaggedTensor、KeyedJaggedTensor 和 KeyedTensor。

class torchrec.sparse.jagged_tensor.JaggedTensor(*args, **kwargs)¶

表示 (可選的加權) 參差不齊的張量。

JaggedTensor 是一個具有參差不齊維度的張量，該維度的切片可能具有不同的長度。有關完整範例，請參閱 KeyedJaggedTensor。

實作是 torch.jit.script-able。

注意

我們將不做輸入驗證，因為它很耗費資源，您應該始終傳入有效的長度、偏移量等。

參數:

values (torch.Tensor) – 密集表示中的值張量。
weights (Optional[torch.Tensor]) – 如果值具有權重。與值具有相同形狀的張量。
lengths (Optional[torch.Tensor]) – 參差不齊的切片，表示為長度。
offsets (Optional[torch.Tensor]) – 參差不齊的切片，表示為累積偏移量。

device() → device¶

取得 JaggedTensor 裝置。

返回:: 值張量的裝置。
返回類型:: torch.device

static empty(is_weighted: bool = False, device: Optional[device] = None, values_dtype: Optional[dtype] = None, weights_dtype: Optional[dtype] = None, lengths_dtype: dtype = torch.int32) → JaggedTensor¶

建構一個空的 JaggedTensor。

參數:

is_weighted (bool) – JaggedTensor 是否有權重。
device (Optional[torch.device]) – JaggedTensor 的裝置。
values_dtype (Optional[torch.dtype]) – 值的 dtype。
weights_dtype (Optional[torch.dtype]) – 權重的 dtype。
lengths_dtype (torch.dtype) – 長度的 dtype。

返回:

空的 JaggedTensor。

返回類型:

JaggedTensor

static from_dense(values: List[Tensor], weights: Optional[List[Tensor]] = None) → JaggedTensor¶

從張量列表建構 JaggedTensor 作為值，可選擇權重。lengths 將被計算，形狀為 (B,)，其中 B 是 len(values)，代表批次大小。

參數:

values (List[torch.Tensor]) – 用於密集表示的張量列表
weights (Optional[List[torch.Tensor]]) – 如果值有權重，則為與值形狀相同的張量。

返回:

從 2D 密集張量建立的 JaggedTensor。

返回類型:

JaggedTensor

範例

values = [
    torch.Tensor([1.0]),
    torch.Tensor(),
    torch.Tensor([7.0, 8.0]),
    torch.Tensor([10.0, 11.0, 12.0]),
]
weights = [
    torch.Tensor([1.0]),
    torch.Tensor(),
    torch.Tensor([7.0, 8.0]),
    torch.Tensor([10.0, 11.0, 12.0]),
]
j1 = JaggedTensor.from_dense(
    values=values,
    weights=weights,
)

# j1 = [[1.0], [], [7.0, 8.0], [10.0, 11.0, 12.0]]

static from_dense_lengths(values: Tensor, lengths: Tensor, weights: Optional[Tensor] = None) → JaggedTensor¶

從值和長度張量建構 JaggedTensor，可選擇權重。請注意，lengths 仍然是形狀 (B,)，其中 B 是批次大小。

參數:

values (torch.Tensor) – 值的密集表示。
lengths (torch.Tensor) – 鋸齒狀切片，表示為長度。
weights (Optional[torch.Tensor]) – 如果值有權重，則為與值形狀相同的張量。

返回:

從 2D 密集張量建立的 JaggedTensor。

返回類型:

JaggedTensor

lengths() → Tensor¶

取得 JaggedTensor 長度。如果未計算，則從偏移量計算。

返回:: 長度張量。
返回類型:: torch.Tensor

lengths_or_none() → Optional[Tensor]¶

取得 JaggedTensor 長度。如果未計算，則傳回 None。

返回:: 長度張量。
返回類型:: Optional[torch.Tensor]

offsets() → Tensor¶

取得 JaggedTensor 偏移量。如果未計算，則從長度計算。

返回:: 偏移量張量。
返回類型:: torch.Tensor

offsets_or_none() → Optional[Tensor]¶

取得 JaggedTensor 的 offsets。如果尚未計算，則傳回 None。

返回:: 偏移量張量。
返回類型:: Optional[torch.Tensor]

record_stream(stream: Stream) → None¶: 請參閱 https://pytorch.dev.org.tw/docs/stable/generated/torch.Tensor.record_stream.html

to(device: device, non_blocking: bool = False) → JaggedTensor¶

將 JaggedTensor 移動到指定的裝置。

參數:

device (torch.device) – 要移動到的裝置。
non_blocking (bool) – 是否執行非同步複製。

返回:

移動後的 JaggedTensor。

返回類型:

JaggedTensor

to_dense() → List[Tensor]¶

建構 JT 值的密集表示法。

返回:: tensor 的列表。
返回類型:: List[torch.Tensor]

範例

values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
offsets = torch.IntTensor([0, 2, 2, 3, 4, 5, 8])
jt = JaggedTensor(values=values, offsets=offsets)

values_list = jt.to_dense()

# values_list = [
#     torch.tensor([1.0, 2.0]),
#     torch.tensor([]),
#     torch.tensor([3.0]),
#     torch.tensor([4.0]),
#     torch.tensor([5.0]),
#     torch.tensor([6.0, 7.0, 8.0]),
# ]

to_dense_weights() → Optional[List[Tensor]]¶

建構 JT 權重的密集表示法。

返回:: tensor 的列表，如果沒有權重則為 None。
返回類型:: Optional[List[torch.Tensor]]

範例

values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
weights = torch.Tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])
offsets = torch.IntTensor([0, 2, 2, 3, 4, 5, 8])
jt = JaggedTensor(values=values, weights=weights, offsets=offsets)

weights_list = jt.to_dense_weights()

# weights_list = [
#     torch.tensor([0.1, 0.2]),
#     torch.tensor([]),
#     torch.tensor([0.3]),
#     torch.tensor([0.4]),
#     torch.tensor([0.5]),
#     torch.tensor([0.6, 0.7, 0.8]),
# ]

to_padded_dense(desired_length: Optional[int] = None, padding_value: float = 0.0) → Tensor¶

從 JT 的值建構形狀為 (B, N,) 的 2D 密集 tensor。

請注意，B 是 self.lengths() 的長度，而 N 是最長的特徵長度或 desired_length。

如果 desired_length > length，我們將使用 padding_value 進行填充，否則我們將在 desired_length 處選擇最後一個值。

參數:

desired_length (int) – tensor 的長度。
padding_value (float) – 如果我們需要填充，則為填充值。

返回:

2d 密集 tensor。

返回類型:

torch.Tensor

範例

values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
offsets = torch.IntTensor([0, 2, 2, 3, 4, 5, 8])
jt = JaggedTensor(values=values, offsets=offsets)

dt = jt.to_padded_dense(
    desired_length=2,
    padding_value=10.0,
)

# dt = [
#     [1.0, 2.0],
#     [10.0, 10.0],
#     [3.0, 10.0],
#     [4.0, 10.0],
#     [5.0, 10.0],
#     [6.0, 7.0],
# ]

to_padded_dense_weights(desired_length: Optional[int] = None, padding_value: float = 0.0) → Optional[Tensor]¶

從 JT 的權重建構形狀為 (B, N,) 的 2D 密集 tensor。

請注意，B （批次大小）是 self.lengths() 的長度，而 N 是最長的特徵長度或 desired_length。

如果 desired_length > length，我們將使用 padding_value 進行填充，否則我們將在 desired_length 處選擇最後一個值。

類似於 to_padded_dense，但適用於 JT 的權重而不是值。

參數:

desired_length (int) – tensor 的長度。
padding_value (float) – 如果我們需要填充，則為填充值。

返回:

2d 密集 tensor，如果沒有權重則為 None。

返回類型:

Optional[torch.Tensor]

範例

values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
weights = torch.Tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])
offsets = torch.IntTensor([0, 2, 2, 3, 4, 5, 8])
jt = JaggedTensor(values=values, weights=weights, offsets=offsets)

d_wt = jt.to_padded_dense_weights(
    desired_length=2,
    padding_value=1.0,
)

# d_wt = [
#     [0.1, 0.2],
#     [1.0, 1.0],
#     [0.3, 1.0],
#     [0.4, 1.0],
#     [0.5, 1.0],
#     [0.6, 0.7],
# ]

values() → Tensor¶

取得 JaggedTensor 值。

返回:: 值 tensor。
返回類型:: torch.Tensor

weights() → Tensor¶

取得 JaggedTensor 權重。如果為 None，則拋出錯誤。

返回:: 權重 tensor。
返回類型:: torch.Tensor

weights_or_none() → Optional[Tensor]¶

取得 JaggedTensor 權重。如果為 None，則傳回 None。

返回:: 權重 tensor。
返回類型:: Optional[torch.Tensor]

class torchrec.sparse.jagged_tensor.KeyedJaggedTensor(*args, **kwargs)¶

表示（可選擇加權的）keyed jagged tensor。

KeyedJaggedTensor 是一個具有jagged dimension的 tensor，該 dimension 是其切片可能具有不同長度的 dimension。在第一維度上 keyed，在最後一維度上 jagged。

實作是 torch.jit.script-able。

參數:

keys (List[str]) – 指向不規則張量的鍵 (keys)。
values (torch.Tensor) – 密集表示中的值張量。
weights (Optional[torch.Tensor]) – 如果數值 (values) 具有權重 (weights)。與數值 (values) 具有相同形狀的張量。
lengths (Optional[torch.Tensor]) – 參差不齊的切片，表示為長度。
offsets (Optional[torch.Tensor]) – 參差不齊的切片，表示為累積偏移量。
stride (Optional[int]) – 每個批次 (batch) 的範例數量。
stride_per_key_per_rank (Optional[List[List[int]]]) – 每個鍵 (key) 每個等級 (rank) 的批次大小 (範例數量)，其中外部列表代表鍵，內部列表代表數值。內部列表中的每個值表示在分散式環境中，來自其索引等級 (rank) 的批次 (batch) 中的範例數量。
length_per_key (Optional[List[int]]) – 每個鍵的起始長度。
offset_per_key (Optional[List[int]]) – 每個鍵的起始偏移量 (offset) 和最終偏移量。
index_per_key (Optional[Dict[str, int]]) – 每個鍵的索引。
jt_dict (Optional[Dict[str, JaggedTensor]]) – 鍵 (keys) 對應到 JaggedTensors 的字典 (dictionary)。允許 to_dict() 具有延遲/可快取的能力。
inverse_indices (Optional[Tuple[List[str], torch.Tensor]]) – 反向索引，用於展開每個鍵 (key) 可變步幅 (stride) 的重複資料刪除嵌入輸出 (deduplicated embedding output)。

範例

#              0       1        2  <-- dim_1
# "Feature0"   [V0,V1] None    [V2]
# "Feature1"   [V3]    [V4]    [V5,V6,V7]
#   ^
#  dim_0

dim_0: keyed dimension (ie. `Feature0`, `Feature1`)
dim_1: optional second dimension (ie. batch size)
dim_2: The jagged dimension which has slice lengths between 0-3 in the above example

# We represent this data with following inputs:

values: torch.Tensor = [V0, V1, V2, V3, V4, V5, V6, V7]  # V == any tensor datatype
weights: torch.Tensor = [W0, W1, W2, W3, W4, W5, W6, W7]  # W == any tensor datatype
lengths: torch.Tensor = [2, 0, 1, 1, 1, 3]  # representing the jagged slice
offsets: torch.Tensor = [0, 2, 2, 3, 4, 5, 8]  # offsets from 0 for each jagged slice
keys: List[str] = ["Feature0", "Feature1"]  # correspond to each value of dim_0
index_per_key: Dict[str, int] = {"Feature0": 0, "Feature1": 1}  # index for each key
offset_per_key: List[int] = [0, 3, 8]  # start offset for each key and final offset

static concat(kjt_list: List[KeyedJaggedTensor]) → KeyedJaggedTensor¶

將 KeyedJaggedTensors 列表串聯成單一 KeyedJaggedTensor。

參數:: kjt_list (List[KeyedJaggedTensor]) – 要串聯的 KeyedJaggedTensors 列表。
返回:: 串聯後的 KeyedJaggedTensor。
返回類型:: KeyedJaggedTensor

device() → device¶

傳回 KeyedJaggedTensor 的 device。

返回:: KeyedJaggedTensor 的 device。
返回類型:: torch.device

static empty(is_weighted: bool = False, device: Optional[device] = None, values_dtype: Optional[dtype] = None, weights_dtype: Optional[dtype] = None, lengths_dtype: dtype = torch.int32) → KeyedJaggedTensor¶

建構一個空的 KeyedJaggedTensor。

參數:

is_weighted (bool) – KeyedJaggedTensor 是否為加權 (weighted)。
device (Optional[torch.device]) – KeyedJaggedTensor 將被放置的 device。
values_dtype (Optional[torch.dtype]) – 數值 (values) 張量的 dtype。
weights_dtype (Optional[torch.dtype]) – 權重 (weights) 張量的 dtype。
lengths_dtype (torch.dtype) – 長度 (lengths) 張量的 dtype。

返回:

空的 KeyedJaggedTensor。

返回類型:

KeyedJaggedTensor

static empty_like(kjt: KeyedJaggedTensor) → KeyedJaggedTensor¶

建構一個空的 KeyedJaggedTensor，其 device 和 dtypes 與輸入的 KeyedJaggedTensor 相同。

參數:: kjt (KeyedJaggedTensor) – 輸入的 KeyedJaggedTensor。
返回:: 空的 KeyedJaggedTensor。
返回類型:: KeyedJaggedTensor

static from_jt_dict(jt_dict: Dict[str, JaggedTensor]) → KeyedJaggedTensor¶

從 JaggedTensors 的字典 (dictionary) 建構 KeyedJaggedTensor。會自動在新建立的 KJT 上呼叫 kjt.sync()。

注意

只有當 JaggedTensors 都具有相同的“隱式” batch_size 維度時，此函數才會起作用。

基本上，我們可以將 JaggedTensor 視覺化為 [batch_size x variable_feature_dim] 格式的二維張量。在某些情況下，如果某個 batch 沒有特徵值，則輸入的 JaggedTensor 可能就不包含任何值。

但 KeyedJaggedTensor (預設情況下) 通常會填充 “None”，以便 KeyedJaggedTensor 中儲存的所有 JaggedTensor 都具有相同的 batch_size 維度。也就是說，如果 JaggedTensor 輸入沒有自動為空的 batch 填充，此函式會出錯/無法運作。

考慮以下 KeyedJaggedTensor 的視覺化：# 0 1 2 <– dim_1 # “Feature0” [V0,V1] None [V2] # “Feature1” [V3] [V4] [V5,V6,V7] # ^ # dim_0

現在，如果輸入 jt_dict = {: # “Feature0” [V0,V1] [V2] # “Feature1” [V3] [V4] [V5,V6,V7]

} 並且每個 JaggedTensor 中都遺漏了 “None”，則此函式將會失敗，因為我們無法正確填充 “None”，因為它在技術上不知道要在 JaggedTensor 內的哪個正確的 batch/位置進行填充。

本質上，此函式推斷的 lengths 張量將為 [2, 1, 1, 1, 3]，表示可變的 batch_size dim_1 違反了現有的假設/前提條件，即 KeyedJaggedTensor 應該具有固定的 batch_size 維度。

參數:: jt_dict (Dict[str, JaggedTensor]) – JaggedTensor 的字典。
返回:: 建構的 KeyedJaggedTensor。
返回類型:: KeyedJaggedTensor

static from_lengths_sync(keys: List[str], values: Tensor, lengths: Tensor, weights: Optional[Tensor] = None, stride: Optional[int] = None, stride_per_key_per_rank: Optional[List[List[int]]] = None, inverse_indices: Optional[Tuple[List[str], Tensor]] = None) → KeyedJaggedTensor¶

從鍵、長度和偏移量的清單建構 KeyedJaggedTensor。與 from_offsets_sync 相同，只是使用長度而不是偏移量。

參數:

keys (List[str]) – 鍵的清單。
values (torch.Tensor) – 密集表示中的值張量。
lengths (torch.Tensor) – 鋸齒狀切片，表示為長度。
weights (Optional[torch.Tensor]) – 如果數值 (values) 具有權重 (weights)。與數值 (values) 具有相同形狀的張量。
stride (Optional[int]) – 每個批次 (batch) 的範例數量。
stride_per_key_per_rank (Optional[List[List[int]]]) – 每個 rank 每個鍵的 batch 大小（範例數量），其中外部清單表示鍵，內部清單表示值。
inverse_indices (Optional[Tuple[List[str], torch.Tensor]]) – 反向索引，用於展開每個鍵 (key) 可變步幅 (stride) 的重複資料刪除嵌入輸出 (deduplicated embedding output)。

返回:

建構的 KeyedJaggedTensor。

返回類型:

KeyedJaggedTensor

static from_offsets_sync(keys: List[str], values: Tensor, offsets: Tensor, weights: Optional[Tensor] = None, stride: Optional[int] = None, stride_per_key_per_rank: Optional[List[List[int]]] = None, inverse_indices: Optional[Tuple[List[str], Tensor]] = None) → KeyedJaggedTensor¶

從鍵、值和偏移量的清單建構 KeyedJaggedTensor。

參數:

keys (List[str]) – 鍵的清單。
values (torch.Tensor) – 密集表示中的值張量。
offsets (torch.Tensor) – Jagged slices，表示為累積偏移量。
weights (Optional[torch.Tensor]) – 如果數值 (values) 具有權重 (weights)。與數值 (values) 具有相同形狀的張量。
stride (Optional[int]) – 每個批次 (batch) 的範例數量。
stride_per_key_per_rank (Optional[List[List[int]]]) – 每個 rank 每個鍵的 batch 大小（範例數量），其中外部清單表示鍵，內部清單表示值。
inverse_indices (Optional[Tuple[List[str], torch.Tensor]]) – 反向索引，用於展開每個鍵 (key) 可變步幅 (stride) 的重複資料刪除嵌入輸出 (deduplicated embedding output)。

返回:

建構的 KeyedJaggedTensor。

返回類型:

KeyedJaggedTensor

index_per_key() → Dict[str, int]¶

傳回 KeyedJaggedTensor 的每個鍵的索引。

返回:: KeyedJaggedTensor 的每個鍵的索引。
返回類型:: Dict[str, int]

inverse_indices() → Tuple[List[str], Tensor]¶

傳回 KeyedJaggedTensor 的反向索引。如果反向索引為 None，則會拋出錯誤。

返回:: KeyedJaggedTensor 的反向索引。
返回類型:: Tuple[List[str], torch.Tensor]

inverse_indices_or_none() → Optional[Tuple[List[str], Tensor]]¶

傳回 KeyedJaggedTensor 的反向索引，如果不存在則傳回 None。

返回:: KeyedJaggedTensor 的反向索引。
返回類型:: Optional[Tuple[List[str], torch.Tensor]]

keys() → List[str]¶

傳回 KeyedJaggedTensor 的鍵 (keys)。

返回:: KeyedJaggedTensor 的鍵 (keys)。
返回類型:: List[str]

length_per_key() → List[int]¶

傳回 KeyedJaggedTensor 每個鍵的長度。如果每個鍵的長度為 None，則會計算它。

返回:: KeyedJaggedTensor 每個鍵的長度。
返回類型:: List[int]

length_per_key_or_none() → Optional[List[int]]¶

傳回 KeyedJaggedTensor 每個鍵的長度，如果尚未計算，則傳回 None。

返回:: KeyedJaggedTensor 每個鍵的長度。
返回類型:: List[int]

lengths() → Tensor¶

傳回 KeyedJaggedTensor 的長度。如果尚未計算長度，將會計算它們。

返回:: KeyedJaggedTensor 的長度。
返回類型:: torch.Tensor

lengths_offset_per_key() → List[int]¶

傳回 KeyedJaggedTensor 每個鍵的長度偏移。如果每個鍵的長度偏移為 None，則會計算它。

返回:: KeyedJaggedTensor 每個鍵的長度偏移。
返回類型:: List[int]

lengths_or_none() → Optional[Tensor]¶

傳回 KeyedJaggedTensor 的長度，如果尚未計算，則傳回 None。

返回:: KeyedJaggedTensor 的長度。
返回類型:: torch.Tensor

offset_per_key() → List[int]¶

傳回 KeyedJaggedTensor 每個鍵的偏移量。如果每個鍵的偏移量為 None，則會計算它。

返回:: KeyedJaggedTensor 每個鍵的偏移量。
返回類型:: List[int]

offset_per_key_or_none() → Optional[List[int]]¶

傳回 KeyedJaggedTensor 每個鍵的偏移量，如果尚未計算，則傳回 None。

返回:: KeyedJaggedTensor 每個鍵的偏移量。
返回類型:: List[int]

offsets() → Tensor¶

傳回 KeyedJaggedTensor 的偏移量。如果尚未計算偏移量，將會計算它們。

返回:: KeyedJaggedTensor 的偏移量。
返回類型:: torch.Tensor

offsets_or_none() → Optional[Tensor]¶

傳回 KeyedJaggedTensor 的偏移量，如果尚未計算，則傳回 None。

返回:: KeyedJaggedTensor 的偏移量。
返回類型:: torch.Tensor

permute(indices: List[int], indices_tensor: Optional[Tensor] = None) → KeyedJaggedTensor¶

排列 KeyedJaggedTensor。

參數:

indices (List[int]) – 索引的列表。
indices_tensor (Optional[torch.Tensor]) – 索引的張量 (tensor)。

返回:

排列後的 KeyedJaggedTensor。

返回類型:

KeyedJaggedTensor

record_stream(stream: Stream) → None¶: 請參閱 https://pytorch.dev.org.tw/docs/stable/generated/torch.Tensor.record_stream.html

split(segments: List[int]) → List[KeyedJaggedTensor]¶

將 KeyedJaggedTensor 分割成 KeyedJaggedTensor 的列表。

參數:: segments (List[int]) – 分割區段的列表。
返回:: KeyedJaggedTensor 的列表。
返回類型:: List[KeyedJaggedTensor]

stride() → int¶

返回 KeyedJaggedTensor 的 stride。如果 stride 為 None，則會計算它。

返回:: KeyedJaggedTensor 的 stride。
返回類型:: int

stride_per_key() → List[int]¶

返回 KeyedJaggedTensor 每個鍵 (key) 的 stride。如果每個鍵的 stride 為 None，則會計算它。

返回:: KeyedJaggedTensor 每個鍵的 stride。
返回類型:: List[int]

stride_per_key_per_rank() → List[List[int]]¶

返回 KeyedJaggedTensor 每個鍵、每個秩 (rank) 的 stride。

返回:: KeyedJaggedTensor 每個鍵、每個秩的 stride。
返回類型:: List[List[int]]

sync() → KeyedJaggedTensor¶

通過計算 offset_per_key 和 length_per_key 來同步 KeyedJaggedTensor。

返回:: 已同步的 KeyedJaggedTensor。
返回類型:: KeyedJaggedTensor

to(device: device, non_blocking: bool = False, dtype: Optional[dtype] = None) → KeyedJaggedTensor¶

返回指定 device 和 dtype 中 KeyedJaggedTensor 的副本。

參數:

device (torch.device) – 副本的目標 device。
non_blocking (bool) – 是否以非阻塞方式複製張量。
dtype (Optional[torch.dtype]) – 副本的目標資料類型。

返回:

複製的 KeyedJaggedTensor。

返回類型:

KeyedJaggedTensor

to_dict() → Dict[str, JaggedTensor]¶

返回每個鍵的 JaggedTensor 字典。將結果緩存在 self._jt_dict 中。

返回:: 每個鍵的 JaggedTensor 字典。
返回類型:: Dict[str, JaggedTensor]

unsync() → KeyedJaggedTensor¶

通過清除 offset_per_key 和 length_per_key 來取消同步 KeyedJaggedTensor。

返回:: 未同步的 KeyedJaggedTensor。
返回類型:: KeyedJaggedTensor

values() → Tensor¶

傳回 KeyedJaggedTensor 的數值。

返回:: KeyedJaggedTensor 的數值。
返回類型:: torch.Tensor

variable_stride_per_key() → bool¶

傳回 KeyedJaggedTensor 是否具有每個鍵的可變步幅。

返回:: KeyedJaggedTensor 是否具有每個鍵的可變步幅。
返回類型:: bool

weights() → Tensor¶

傳回 KeyedJaggedTensor 的權重。如果權重為 None，則會拋出錯誤。

返回:: KeyedJaggedTensor 的權重。
返回類型:: torch.Tensor

weights_or_none() → Optional[Tensor]¶

傳回 KeyedJaggedTensor 的權重，如果不存在則傳回 None。

返回:: KeyedJaggedTensor 的權重。
返回類型:: torch.Tensor

class torchrec.sparse.jagged_tensor.KeyedTensor(*args, **kwargs)¶

KeyedTensor 保存一個串聯的密集張量列表，每個張量都可以通過一個鍵來訪問。

鍵控維度可以是可變長度 (length_per_key)。常見的使用案例包括儲存不同維度的池化嵌入。

實作是 torch.jit.script-able。

參數:

keys (List[str]) – 鍵的清單。
length_per_key (List[int]) – 每個鍵沿著鍵維度的長度。
values (torch.Tensor) – 密集張量，通常沿著鍵維度串聯。
key_dim (int) – 鍵維度，從零開始索引 - 預設為 1 (通常 B 是 0 維度)。

範例

# kt is KeyedTensor holding

#                         0           1           2
#     "Embedding A"    [1,1]       [1,1]        [1,1]
#     "Embedding B"    [2,1,2]     [2,1,2]      [2,1,2]
#     "Embedding C"    [3,1,2,3]   [3,1,2,3]    [3,1,2,3]

tensor_list = [
    torch.tensor([[1,1]] * 3),
    torch.tensor([[2,1,2]] * 3),
    torch.tensor([[3,1,2,3]] * 3),
]

keys = ["Embedding A", "Embedding B", "Embedding C"]

kt = KeyedTensor.from_tensor_list(keys, tensor_list)

kt.values()
# torch.Tensor(
#     [
#         [1, 1, 2, 1, 2, 3, 1, 2, 3],
#         [1, 1, 2, 1, 2, 3, 1, 2, 3],
#         [1, 1, 2, 1, 2, 3, 1, 2, 3],
#     ]
# )

kt["Embedding B"]
# torch.Tensor([[2, 1, 2], [2, 1, 2], [2, 1, 2]])

device() → device¶

返回:: values 張量的裝置。
返回類型:: torch.device

static from_tensor_list(keys: List[str], tensors: List[Tensor], key_dim: int = 1, cat_dim: int = 1) → KeyedTensor¶

從張量列表建立 KeyedTensor。張量沿著 cat_dim 串聯。鍵用於索引張量。

參數:

keys (List[str]) – 鍵的清單。
tensors (List[torch.Tensor]) – 張量列表。
key_dim (int) – 鍵維度，從零開始索引 - 預設為 1 (通常 B 是 0 維度)。
cat_dim (int) – 沿著哪個維度串聯張量 - 預設為

返回:

鍵控張量。

返回類型:

KeyedTensor

key_dim() → int¶

返回:: 鍵維度，從零開始索引 - 通常 B 是 0 維度。
返回類型:: int

keys() → List[str]¶

返回:: 鍵的列表。
返回類型:: List[str]

length_per_key() → List[int]¶

返回:: 每個鍵沿著鍵維度的長度。
返回類型:: List[int]

offset_per_key() → List[int]¶

獲取每個鍵沿著鍵維度的偏移量。如果尚未計算，則計算並快取。

返回:: 每個鍵沿著鍵維度的偏移量。
返回類型:: List[int]

record_stream(stream: Stream) → None¶: 請參閱 https://pytorch.dev.org.tw/docs/stable/generated/torch.Tensor.record_stream.html

static regroup(keyed_tensors: List[KeyedTensor], groups: List[List[str]]) → List[Tensor]¶

將 KeyedTensor 的列表重新分組為 tensor 的列表。

參數:

keyed_tensors (List[KeyedTensor]) – KeyedTensor 的列表。
groups (List[List[str]]) – keys 的群組列表。

返回:

tensor 的列表。

返回類型:

List[torch.Tensor]

static regroup_as_dict(keyed_tensors: List[KeyedTensor], groups: List[List[str]], keys: List[str]) → Dict[str, Tensor]¶

將 KeyedTensor 的列表重新分組為 tensor 的字典。

參數:

keyed_tensors (List[KeyedTensor]) – KeyedTensor 的列表。
groups (List[List[str]]) – keys 的群組列表。
keys (List[str]) – 鍵的清單。

返回:

tensor 的字典。

返回類型:

Dict[str, torch.Tensor]

to(device: device, non_blocking: bool = False) → KeyedTensor¶

將 values tensor 移動到指定的設備。

參數:

device (torch.device) – 要將 values tensor 移動到的設備。
non_blocking (bool) – 是否以非同步方式執行操作 (預設：False)。

返回:

Keyed tensor，其 values tensor 已移動到指定的設備。

返回類型:

KeyedTensor

to_dict() → Dict[str, Tensor]¶

返回:: 由 keys 鍵控的 tensor 字典。
返回類型:: Dict[str, torch.Tensor]

values() → Tensor¶

取得 values tensor。

返回:: 密集 tensor，通常沿著 key 維度串連。
返回類型:: torch.Tensor

資料類型¶

文件

Tutorials

Resources