企业网站建站模板深圳市建筑人才网

TorchRec中的输入和输出格式

---

---

前言

• TorchRec具有其特定的输入输出格式，跟Torch中的Tonsor有些不同，下边就让我们来了解他们。

---

一、JaggedTensor

• JaggedTensor 通过长度、值和偏移量来表示稀疏特征。它之所以被称为“jagged”，是因为它可以有效地表示可变长度序列的数据。
• 相比之下，规范的 torch.Tensor 假设每个序列具有相同的长度，但这在真实世界数据中通常不是这种情况。
• JaggedTensor 有助于表示此类数据而无需填充，从而使其非常高效。

1.1 核心概念

• JaggedTensor表示一个包含多个不等长序列的2D张量，例如：
  • 用户历史点击序列：[[item1, item2], [item3], [item4, item5, item6], …]
  • 文本分词后的句子：[[tokenA, tokenB], [tokenC], …]

1.2 核心属性，也就是参数

• values (Tensor): 一个 1D 张量，包含每个实体的实际值，连续存储，形状为[total_values]
• lengths (Optional[Tensor]): 一个整数列表，表示每个实体的元素数量，形状为[B]
• offsets (Optional[Tensor]): 一个整数列表，表示扁平化值张量中每个序列的起始索引。这些提供了长度的替代方案。形状为[B+1]

演示代码如下：

import torch
from torchrec.sparse.jagged_tensor import JaggedTensor# 方式1 - 使用lengths
values = torch.tensor([1, 2, 3, 4, 5])
lengths = torch.tensor([2, 1, 2])  # 三个序列的长度分别为2,1,2
jt = JaggedTensor(lengths=lengths, values=values)
# 意思就是有三个张量[[1, 2],[3],[4, 5]]# 方式2 - 使用offsets
values = torch.tensor([10, 20, 30, 40])
offsets = torch.tensor([1, 3])  # 两个序列的偏移量
jt = JaggedTensor(offsets=offsets, values=values)
# 意思就是有两个张量[[10],[20, 30, 40]]

1.3 关键操作与方法

• 操作：
  • 拼接

  torchrec.sparse.jagged_ops.concat(jt1, jt2)
  

  • 分块

  jt.split(split_size)
  

  • 聚合

  jt.sum(dim=1) 或 jt.mean(dim=1)
  

• 方法
  • 与稠密张量互转（）

  # 转Padded Tensor
  padded, mask = jt.to_padded_tensor(padding_value=0, max_length=5)
  # 转Packed Tensor (类似PyTorch的PackedSequence)
  packed = jt.to_packed_tensor()
  

  • 嵌入表查询

  # 转Padded Tensor
  embedding_bag = torch.nn.EmbeddingBag(num_embeddings=100, embedding_dim=16)
  embeddings = embedding_bag(jt.values, offsets=jt.offsets())
  

二、KeyedJaggedTensor

• KeyedJaggedTensor 通过引入键（通常是特征名称）来标记不同的特征组（例如，用户特征和项目特征），从而扩展了 JaggedTensor 的功能。
• 这是 EmbeddingBagCollection 和 EmbeddingCollection 的 forward 中使用的数据类型，因为它们用于表示表中的多个特征。

2.1 核心概念

• KeyedJaggedTensor 是用于管理多个变长特征序列的高效数据结构，核心场景包括：

  • 多特征推荐系统：同时处理用户历史点击（click_ids）、搜索词（search_ids）、收藏商品（favor_ids）等不同特征

2.2 核心属性，也就是参数

• keys (List[str]): 特征名称列表，如 [“click”, “search”]
• values (Tensor): 一个 1D 张量，所有特征值的展平拼接，连续存储，形状为[total_values]
• lengths (Optional[Tensor]): 一个整数列表，表示每个特征在样本中的长度，按 keys 顺序排列
• offsets (Optional[Tensor]): 一个整数列表，表示扁平化值张量中每个特征列的起始索引。形状为[B+1]

演示代码如下：

import torch
from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, JaggedTensor# 方式一：从多个JaggedTensor构建
# 创建两个特征的JaggedTensor
click_jt = JaggedTensor(lengths=torch.tensor([2, 1]),values=torch.tensor([101, 202, 303])
)
search_jt = JaggedTensor(lengths=torch.tensor([1, 3]),values=torch.tensor([401, 402, 403, 404])
)# 合并为KeyedJaggedTensor
kjt = KeyedJaggedTensor.from_jagged_tensors(keys=["click", "search"],tensors=[click_jt, search_jt]
)print(kjt)
# KeyedJaggedTensor({
#    "click": JaggedTensor([[101, 202], [303]]),
#    "search": JaggedTensor([[401], [402, 403, 404]])
# })# 方式二：从原始数据直接构建
kjt = KeyedJaggedTensor(keys=["click", "search"],values=torch.tensor([101, 202, 303, 401, 402, 403, 404]),lengths=torch.tensor([2, 1, 1, 3]),  # click_lengths + search_lengthsoffsets=None  # 自动生成
)
# 跟上边是一样的效果

3、KeyedTensor

• torch.Tensor 的包装器，允许通过键访问张量值。

---

总结

• 本节我们学习TorchRec中的数据类型，了解他的输入输出格式。