一、TorchRec里边的输入输出类型

TorchRec中的输入和输出格式

文章目录

TorchRec中的输入和输出格式
前言
一、JaggedTensor
- 1.1 核心概念
- 1.2 核心属性，也就是参数
- 1.3 关键操作与方法
二、KeyedJaggedTensor
- 2.1 核心概念
- 2.2 核心属性，也就是参数
3、KeyedTensor
总结

前言

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  # 自动生成
)
# 跟上边是一样的效果