数据类型¶

TorchRec 包含表示嵌入的数据类型，也称为稀疏特征。稀疏特征通常是用于喂入嵌入表的索引。对于给定的一批数据，嵌入查找索引的数量是可变的。因此，需要一个 jagged 维度来表示一批数据中可变数量的嵌入查找索引。

这个部分涵盖了用于表示稀疏特征的3个TorchRec数据类型类： JaggedTensor，KeyedJaggedTensor和KeyedTensor。

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

表示一个（可选加权）锯齿形张量。

一个 JaggedTensor 是一个具有 不规则维度 的张量，其切片可能有不同的长度。请参见 KeyedJaggedTensor 获取完整示例。

实现是可导出为torch.jit.script的。

注意

我们不会进行输入验证，因为它很昂贵，你应该始终传递有效的长度、偏移量等。

Parameters:

值 (torch.Tensor) – 值张量在密集表示中。
权重 (可选[torch.Tensor]) – 如果值有权重。具有相同形状的张量。
长度 (可选[torch.Tensor]) – 指定的切片，以长度表示。
偏移量 (可选[torch.Tensor]) – 表示累积偏移的跳跃切片。

device() → device¶

获取JaggedTensor设备。

Returns:: 值张量的设备。
Return type:: 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。

Parameters:

是加权 (bool) – 是否JaggedTensor有权重。
设备 (可选[torch.device]) – JaggedTensor 的设备。
values_dtype (Optional[torch.dtype]) – 值的dtype。
weights_dtype (Optional[torch.dtype]) – 数据类型为权重。
lengths_dtype (torch.dtype) – dtype for lengths.

Returns:

空的 jagged tensor。

Return type:

JaggedTensor

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

构造 JaggedTensor 从张量列表作为值，带有可选权重。 lengths 将被计算，形状为 (B,)，其中 B 是 len(values)，代表批量大小。

Parameters:

值 (List[torch.Tensor]) – 一个包含用于密集表示的张量列表
权重 (可选[列表[torch.Tensor]]) – 如果值有权重，具有与值相同形状的张量。

Returns:

由2D稠密张量创建的JaggedTensor。

Return type:

JaggedTensor

Example:

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 是批量大小。

Parameters:

值 (torch.Tensor) – 稀疏表示的值。
长度 (torch.Tensor) – 挖掘切片，表示为长度。
权重 (可选[torch.Tensor]) – 如果值有权重，具有与值相同形状的张量。

Returns:

由2D稠密张量创建的JaggedTensor。

Return type:

JaggedTensor

lengths() → Tensor¶

获取锯齿形张量的长度。如果没有计算过，从偏移量中计算它。

Returns:: 长度张量。
Return type:: torch.Tensor

lengths_or_none() → Optional[Tensor]¶

获取锯齿张量的长度。如果没有计算过，返回None。

Returns:: 长度张量。
Return type:: Optional[torch.Tensor]

offsets() → Tensor¶

获取锯齿形张量偏移。如果没有计算过，从长度中计算它。

Returns:: 偏移张量。
Return type:: torch.Tensor

offsets_or_none() → Optional[Tensor]¶

获取锯齿形张量偏移。如果没有计算，返回None。

Returns:: 偏移张量。
Return type:: Optional[torch.Tensor]

record_stream(stream: Stream) → None¶: 查看 https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html

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

将抖动张量移动到指定设备。

Parameters:

设备 (torch.device) – 要移动到的设备。
非阻塞 (bool) – 是否异步执行复制操作。

Returns:

移动的JaggedTensor。

Return type:

JaggedTensor

to_dense() → List[Tensor]¶

构建JT的密集表示。

Returns:: 列表中的张量。
Return type:: List[torch.Tensor]

Example:

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的权重密集表示。

Returns:: 列表中的张量，None 如果没有权重。
Return type:: Optional[List[torch.Tensor]]

Example:

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¶

构建一个形状为 (B, N,) 的 2D 稠密张量，其值来自 JT。

注意，B是self.lengths()的长度，而 N是最长特征长度或desired_length。

如果 desired_length 大于 length，我们将填充为 padding_value，否则我们将选择最后一个值在 desired_length。

Parameters:

desired_length (int) – 期望张量的长度。
填充值 (浮点数) – 如果需要填充，则填充值。

Returns:

二维密集张量。

Return type:

torch.Tensor

Example:

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]¶

构建一个形状为 (B, N,) 的 2D 稠密张量，其中的权重来自 JT。

注意，B（批量大小）是self.lengths()的长度， N是最长特征长度或desired_length。

如果 desired_length 大于 length，我们将填充为 padding_value，否则我们将选择最后一个值在 desired_length。

喜欢 to_padded_dense 但为了 JT 的权重而不是值。

Parameters:

desired_length (int) – 期望张量的长度。
填充值 (浮点数) – 如果需要填充，则填充值。

Returns:

2d密集张量，None如果无权重。

Return type:

Optional[torch.Tensor]

Example:

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¶

获取锯齿形张量值。

Returns:: 值张量。
Return type:: torch.Tensor

weights() → Tensor¶

获取JaggedTensor权重。如果为None，则抛出错误。

Returns:: 权重张量。
Return type:: torch.Tensor

weights_or_none() → Optional[Tensor]¶

获取JaggedTensor权重。如果为None，则返回None。

Returns:: 权重张量。
Return type:: Optional[torch.Tensor]

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

表示一个（可选加权）键值的锯齿形张量。

一个 KeyedJaggedTensor 是一个具有不规则维度的张量，其维度中的切片可能有不同的长度。以第一维为键和最后一维为不规则。

实现是可导出为torch.jit.script的。

Parameters:

键 (List[字符串]) – 指向 jagged Tensor 的键。
值 (torch.Tensor) – 值张量在密集表示中。
权重 (可选[torch.Tensor]) – 如果值有权重。具有与值相同形状的张量。
长度 (可选[torch.Tensor]) – 指定的切片，以长度表示。
偏移量 (可选[torch.Tensor]) – 表示累积偏移的跳跃切片。
步长 (可选[整数]) – 批次中的示例数量。
步长每键每秩 (Optional[List[List[整数]]]) – 批次大小（示例数量）每键每秩，其中外列表表示键，内列表表示值。每个内列表中的值表示分布式上下文中其索引的排名批次中示例的数量。
长度每键 (可选[列表[整数]]) – 每个键的起始长度。
偏移量每键 (可选[列表[整数]]) – 每个键的起始偏移量和最终偏移量。
index_per_key (Optional[Dict[str, int]]) – 每个键的索引。
jt_dict (Optional[Dict[str, JaggedTensor]]) – 字典，键为JaggedTensors。允许将to_dict()设置为懒加载/缓存化。
逆向索引 (可选[元组[列表[字符串], torch.Tensor]]) – 用于为每个键设置变量步长时，扩展去重嵌入输出的逆向索引。

Example:

#              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¶

将一系列键值稀疏张量拼接成一个单一的键值稀疏张量。

Parameters:: kjt_list (List[KeyedJaggedTensor]) – 列表中的 KeyedJaggedTensors，用于连接。
Returns:: 拼接的键式JaggedTensor。
Return type:: KeyedJaggedTensor

device() → device¶

返回 KeyedJaggedTensor 的设备。

Returns:: KeyedJaggedTensor 的设备。
Return type:: 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。

Parameters:

是加权 (bool) – 是否为KeyedJaggedTensor加权。
设备 (可选[torch.device]) – KeyedJaggedTensor 将在该设备上放置。
values_dtype (Optional[torch.dtype]) – 值张量的dtype。
weights_dtype (Optional[torch.dtype]) – weights tensor 的数据类型。
lengths_dtype (torch.dtype) – dtype of the lengths tensor.

Returns:

空的键值张量。

Return type:

KeyedJaggedTensor

static empty_like(kjt: KeyedJaggedTensor) → KeyedJaggedTensor¶

构建一个与输入的KeyedJaggedTensor具有相同设备和数据类型的空KeyedJaggedTensor。

Parameters:: kjt (KeyedJaggedTensor) – 输入 KeyedJaggedTensor。
Returns:: 空的键值张量。
Return type:: KeyedJaggedTensor

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

从字典中的JaggedTensors构造一个KeyedJaggedTensor。自动调用kjt.sync()在新创建的KJT上。

注意

这个函数只有在所有JaggedTensors都具有相同的“隐含”批次大小维度时才会工作。

基本上，我们可以将JaggedTensors可视化为具有格式 [batch_size x variable_feature_dim] 的2-D张量。在这种情况下，我们有一些没有特征值的批次，输入的JaggedTensor可以不包含任何值。

但是默认情况下，KeyedJaggedTensor 会填充“None”，使得所有存储在 KeyedJaggedTensor 中的 JaggedTensors 都具有相同的 batch_size 维度。也就是说，在这种情况中，如果输入的 JaggedTensor 没有自动填充空批次，那么这个函数就会出错或无法工作。

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

Now if the input jt_dict = {: # “Feature0” [V0,V1] [V2] # “Feature1” [V3] [V4] [V5,V6,V7]

} 并且每个 JaggedTensor 中的“None”都被省略了，那么这个函数就会失败，因为我们无法正确地填充“None”，因为它实际上不知道在 JaggedTensor 中正确的批次或位置进行填充。

基本上，这个函数推断的张量长度为 [2, 1, 1, 1, 3]，这表明变量 batch_size 维度违反了现有的假设/前提，即 KeyedJaggedTensor 应该有固定的 batch_size 维度。

Parameters:: jt_dict (Dict[str, JaggedTensor]) – 字典中的 JaggedTensors。
Returns:: 构建键值稀疏张量。
Return type:: 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相同，但使用长度代替偏移量。

Parameters:

键 (List[字符串]) – 列表。
值 (torch.Tensor) – 值张量在密集表示中。
长度 (torch.Tensor) – 挖掘切片，表示为长度。
权重 (可选[torch.Tensor]) – 如果值有权重。具有与值相同形状的张量。
步长 (可选[整数]) – 批次中的示例数量。
步长每键每秩 (Optional[List[List[整数]]]) – 批次大小（示例数量）每键每秩，其中外层列表表示键，内层列表表示值。
逆向索引 (可选[元组[列表[字符串], torch.Tensor]]) – 用于为每个键设置变量步长时，扩展去重嵌入输出的逆向索引。

Returns:

构建键值稀疏张量。

Return type:

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。

Parameters:

键 (List[字符串]) – 列表。
值 (torch.Tensor) – 值张量在密集表示中。
偏移量 (torch.Tensor) – 指示切片的累积偏移。
权重 (可选[torch.Tensor]) – 如果值有权重。具有与值相同形状的张量。
步长 (可选[整数]) – 批次中的示例数量。
步长每键每秩 (Optional[List[List[整数]]]) – 批次大小（示例数量）每键每秩，其中外层列表表示键，内层列表表示值。
逆向索引 (可选[元组[列表[字符串], torch.Tensor]]) – 用于为每个键设置变量步长时，扩展去重嵌入输出的逆向索引。

Returns:

构建键值稀疏张量。

Return type:

KeyedJaggedTensor

index_per_key() → Dict[str, int]¶

返回 KeyedJaggedTensor 中每个键的索引。

Returns:: 键值对的KeyedJaggedTensor中的索引。
Return type:: 字典

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

返回KeyedJaggedTensor的逆索引。如果逆索引为None，这将抛出错误。

Returns:: KeyedJaggedTensor 的逆索引。
Return type:: 元组列表[str]，torch.Tensor

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

返回 KeyedJaggedTensor 的逆索引，如果没有，则返回 None。

Returns:: KeyedJaggedTensor 的逆索引。
Return type:: 可选的[Tuple[List[str], torch.Tensor]]

keys() → List[str]¶

返回 KeyedJaggedTensor 的键。

Returns:: KeyedJaggedTensor 的键。
Return type:: List[str]

length_per_key() → List[int]¶

返回 KeyedJaggedTensor 中每个键的长度。如果每个键的长度为 None，则会计算它。

Returns:: 键的KeyedJaggedTensor中的每个键的长度。
Return type:: List[int]

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

返回 KeyedJaggedTensor 的每个键的长度，如果没有计算，则返回 None。

Returns:: 键的KeyedJaggedTensor中的每个键的长度。
Return type:: List[int]

lengths() → Tensor¶

返回 KeyedJaggedTensor 的长度。如果尚未计算这些长度，它将进行计算。

Returns:: KeyedJaggedTensor的长度。
Return type:: torch.Tensor

lengths_offset_per_key() → List[int]¶

返回 KeyedJaggedTensor 中每个键的长度偏移量。如果每个键的长度偏移量为 None，则会计算它。

Returns:: 长度偏移量，每个键的KeyedJaggedTensor。
Return type:: List[int]

lengths_or_none() → Optional[Tensor]¶

返回 KeyedJaggedTensor 的长度，如果没有计算，则返回 None。

Returns:: KeyedJaggedTensor的长度。
Return type:: torch.Tensor

offset_per_key() → List[int]¶

返回KeyedJaggedTensor中每个键的偏移量。如果每个键的偏移量为None，则会计算它。

Returns:: 键的KeyedJaggedTensor中的偏移量。
Return type:: List[int]

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

返回 KeyedJaggedTensor 的每个键的偏移量，如果没有计算，则返回 None。

Returns:: 键的KeyedJaggedTensor中的偏移量。
Return type:: List[int]

offsets() → Tensor¶

返回 KeyedJaggedTensor 的偏移量。如果偏移量尚未计算，它将进行计算。

Returns:: KeyedJaggedTensor 的偏移量。
Return type:: torch.Tensor

offsets_or_none() → Optional[Tensor]¶

返回KeyedJaggedTensor的偏移量，如果没有计算，则返回None。

Returns:: KeyedJaggedTensor 的偏移量。
Return type:: torch.Tensor

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

将键式稀疏张量进行置换。

Parameters:

索引 (列表[整数]) – 列表。
indices_tensor (Optional[torch.Tensor]) – tensor of indices.

Returns:

打乱的键式JaggedTensor。

Return type:

KeyedJaggedTensor

record_stream(stream: Stream) → None¶: 查看 https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html

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

将KeyedJaggedTensor拆分为一个KeyedJaggedTensor列表。

Parameters:: 段落 (List[整数]) – 列表中的段落。
Returns:: KeyedJaggedTensor 的列表。
Return type:: List[KeyedJaggedTensor]

stride() → int¶

返回 KeyedJaggedTensor 的步长。如果步长为 None，则会计算它。

Returns:: KeyedJaggedTensor 的步长。
Return type:: 整数

stride_per_key() → List[int]¶

返回键式稀疏张量的每个键的步长。如果每个键的步长为None，则会计算它。

Returns:: KeyedJaggedTensor 的每个键的步长。
Return type:: List[int]

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

返回键值对的每个键在每个分组中的步长。

Returns:: KeyedJaggedTensor的每个键的每个排名的步长。
Return type:: List[List[int]]

sync() → KeyedJaggedTensor¶

通过计算每个键的偏移量和长度，同步KeyedJaggedTensor。

Returns:: 同步键式JaggedTensor。
Return type:: KeyedJaggedTensor

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

返回指定设备和数据类型的一个副本的KeyedJaggedTensor。

Parameters:

设备 (torch.device) – 复制所需的设备。
非阻塞 (bool) – 是否以非阻塞的方式复制张量。
dtype (Optional[torch.dtype]) – 指定复制的预期数据类型。

Returns:

复制的KeyedJaggedTensor。

Return type:

KeyedJaggedTensor

to_dict() → Dict[str, JaggedTensor]¶

返回每个键的JaggedTensor字典。将结果缓存到self._jt_dict中。

Returns:: 每个键的JaggedTensor字典。
Return type:: Dict[str, JaggedTensor]

unsync() → KeyedJaggedTensor¶

将KeyedJaggedTensor解键化并清除offset_per_key和length_per_key。

Returns:: 未同步的键式跳跃张量。
Return type:: KeyedJaggedTensor

values() → Tensor¶

返回 KeyedJaggedTensor 的值。

Returns:: KeyedJaggedTensor 的值。
Return type:: torch.Tensor

variable_stride_per_key() → bool¶

返回 KeyedJaggedTensor 是否每个键都有可变步长。

Returns:: 是否KeyedJaggedTensor每个键的步长可变。
Return type:: 布尔

weights() → Tensor¶

返回 KeyedJaggedTensor 的权重。如果 weights 为 None，这将抛出错误。

Returns:: KeyedJaggedTensor 的权重。
Return type:: torch.Tensor

weights_or_none() → Optional[Tensor]¶

返回 KeyedJaggedTensor 的权重，如果不存在则返回 None。

Returns:: KeyedJaggedTensor 的权重。
Return type:: torch.Tensor

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

KeyedTensor 持有一个连接的密集张量列表，每个张量都可以通过键访问。

键的维度可以是可变长度（每个键的长度）。常见的使用案例包括存储不同维度的池化嵌入。

实现是可导出为torch.jit.script的。

Parameters:

键 (List[字符串]) – 列表。
长度每键 (List[int]) – 每个键沿键维度的长度。
值 (torch.Tensor) – 密集张量，通常沿着关键维度拼接。
key_dim (int) – key dimension, zero indexed - defaults to 1 (typically B is 0-dimension).

Example:

# 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¶

Returns:: 值张量的设备。
Return type:: torch.device

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

创建一个键值张量列表。这些张量沿着cat_dim进行连接。键用于索引张量。

Parameters:

键 (List[字符串]) – 列表。
张量 (List[PyTorch张量]) – 列表中的张量。
key_dim (int) – key dimension, zero indexed - defaults to 1 (typically B is 0-dimension).
cat_dim (int) – 维度沿着该维度将张量连接 - 默认值

Returns:

键对齐张量。

Return type:

KeyedTensor

key_dim() → int¶

Returns:: 关键维度，零索引 - 通常 B 是 0 维度。
Return type:: 整数

keys() → List[str]¶

Returns:: 列表中的键。
Return type:: List[str]

length_per_key() → List[int]¶

Returns:: 每个键在键维度上的长度。
Return type:: List[int]

offset_per_key() → List[int]¶

获取每个键在键维度上的偏移量。如果尚未计算，则进行计算并缓存。

Returns:: 每个键在键维度上的偏移量。
Return type:: List[int]

record_stream(stream: Stream) → None¶: 查看 https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html

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

将KeyedTensors列表重新组合成一个Tensor列表。

Parameters:

keyed_tensors (List[KeyedTensor]) – 列表中的 KeyedTensors。
群组 (List[List[字符串]]) – 列表中的键的组。

Returns:

列表中的张量。

Return type:

List[torch.Tensor]

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

将一个键值对列表重组为字典形式的张量。

Parameters:

keyed_tensors (List[KeyedTensor]) – 列表中的 KeyedTensors。
群组 (List[List[字符串]]) – 列表中的键的组。
键 (List[字符串]) – 列表。

Returns:

张量字典。

Return type:

字典[str, torch.Tensor]

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

将值张量移动到指定设备。

Parameters:

设备 (torch.device) – 将值张量移动到的设备。
非阻塞 (bool) – 是否异步执行操作（默认：False）。

Returns:

具有键的张量，其值被移动到指定设备。

Return type:

KeyedTensor

to_dict() → Dict[str, Tensor]¶

Returns:: 字典，键为张量的键。
Return type:: 字典[str, torch.Tensor]

values() → Tensor¶

获取值张量。

Returns:: 密集张量，通常沿着关键维度进行连接。
Return type:: torch.Tensor

数据类型¶

文档

教程

资源