TorchRec常用函数方法及嵌入表详解
五、常用的函数方法
文章目录
前言
- 这节我们了解一下常用的函数方法，以及跟之前的pytorch中的函数使用上的一些不同
一、嵌入表
- 嵌入表示为嵌入表中的单独行，也称为嵌入权重。这是因为嵌入/嵌入表权重是通过梯度下降训练的，就像模型的所有其他权重一样！
- 嵌入表只是一个用于存储嵌入的大矩阵，具有二维（B，N），其中：
- B是表存储的嵌入数
- N是每次嵌入的维数（N维嵌入）。
1.1 嵌入的过程
- 嵌入在RecSys中通过以下过程进行训练：
- 1、输入/查找索引作为唯一ID输入到模型中。
- ID被散列到嵌入表的总大小，以防止ID>最大行时出现问题(索引越界)
- 2、检索嵌入并合并，例如取嵌入的总和或均值。
- 这是必需的，因为每个示例可以有一个变量#的嵌入，而模型需要一致的形状。
- 3、嵌入与模型的其余部分结合使用，以产生预测
- 将嵌入好的向量送给不同的模型结构(例如：双塔模型、DIRM等等)，通过模型的输出，我们来预测出一个值，例如：点击率(CTR)
- 4、损失是通过预测和示例的标签计算的
- 模型的所有权重都通过梯度下降和反向传播进行更新，包括与示例相关的嵌入权重。
- 1、输入/查找索引作为唯一ID输入到模型中。
1.2 在torch中的嵌入表
演示代码：
import torch
# 定义嵌入表的大小， num_embeddings是一共有多少个类别、embedding_dim是我们要把每一个唯一的类别映射成多少维的向量表示
num_embeddings, embedding_dim = 10, 4
# 初始化嵌入表
weights = torch.rand(num_embeddings, embedding_dim)
print("Weights:", weights)
# Weights: tensor([[0.3446, 0.3614, 0.8938, 0.8157],
# [0.1830, 0.0326, 0.8241, 0.2995],
# [0.7328, 0.0531, 0.9528, 0.0592],
# [0.7800, 0.1797, 0.0167, 0.7401],
# [0.4837, 0.2052, 0.3360, 0.9656],
# [0.7887, 0.3066, 0.0956, 0.3344],
# [0.5904, 0.8541, 0.5963, 0.2800],
# [0.5751, 0.4341, 0.6218, 0.4101],
# [0.6881, 0.5363, 0.4747, 0.2301],
# [0.6088, 0.1060, 0.1100, 0.7290]])
# 传入我们设定好的权重
embedding_collection = torch.nn.Embedding(
num_embeddings, embedding_dim, _weight=weights
)
embedding_bag_collection = torch.nn.EmbeddingBag(
num_embeddings, embedding_dim, _weight=weights
)
# 打印这两个嵌入层， 应该跟上边的嵌入表一样
print("Embedding Collection Table: ", embedding_collection.weight)
print("Embedding Bag Collection Table: ", embedding_bag_collection.weight)
"""
打印信息 如下：
Embedding Collection Table: Parameter containing:
tensor([[0.3446, 0.3614, 0.8938, 0.8157],
[0.1830, 0.0326, 0.8241, 0.2995],
[0.7328, 0.0531, 0.9528, 0.0592],
[0.7800, 0.1797, 0.0167, 0.7401],
[0.4837, 0.2052, 0.3360, 0.9656],
[0.7887, 0.3066, 0.0956, 0.3344],
[0.5904, 0.8541, 0.5963, 0.2800],
[0.5751, 0.4341, 0.6218, 0.4101],
[0.6881, 0.5363, 0.4747, 0.2301],
[0.6088, 0.1060, 0.1100, 0.7290]], requires_grad=True)
Embedding Bag Collection Table: Parameter containing:
tensor([[0.3446, 0.3614, 0.8938, 0.8157],
[0.1830, 0.0326, 0.8241, 0.2995],
[0.7328, 0.0531, 0.9528, 0.0592],
[0.7800, 0.1797, 0.0167, 0.7401],
[0.4837, 0.2052, 0.3360, 0.9656],
[0.7887, 0.3066, 0.0956, 0.3344],
[0.5904, 0.8541, 0.5963, 0.2800],
[0.5751, 0.4341, 0.6218, 0.4101],
[0.6881, 0.5363, 0.4747, 0.2301],
[0.6088, 0.1060, 0.1100, 0.7290]], requires_grad=True)
"""
代码演示：nn.Embedding和nn.EmbeddingBag的区别
ids = torch.tensor([[1, 3]])
print("Input row IDS: ", ids)
embeddings = embedding_collection(ids)
# 打印嵌入结果，应该看到跟上边对应行的嵌入向量
print("Embedding Collection Results: ")
print(embeddings)
print("Shape: ", embeddings.shape)
"""
打印信息如下
Input row IDS: tensor([[1, 3]])
Embedding Collection Results:
tensor([[[0.1830, 0.0326, 0.8241, 0.2995],
[0.7800, 0.1797, 0.0167, 0.7401]]], grad_fn=<EmbeddingBackward0>)
Shape: torch.Size([1, 2, 4])
"""
# nn.EmbeddedBag 默认池是平均值，因此应该是上述值的批量维度的平均值
pooled_embeddings = embedding_bag_collection(ids)
print("Embedding Bag Collection Results: ")
print(pooled_embeddings)
print("Shape: ", pooled_embeddings.shape)
# nn.EmbeddingBag跟nn.Embedding的结构类似，但是nn.EmbeddingBag会有池化(均值，求和等等)
# 我们可以看到embedding_collection的嵌入平均值和embedding_bag_collection的输出是一样的
print("Mean: ", torch.mean(embedding_collection(ids), dim=1))
"""
打印台输出结果如下：
Embedding Bag Collection Results:
tensor([[0.4815, 0.1062, 0.4204, 0.5198]], grad_fn=<EmbeddingBagBackward0>)
Shape: torch.Size([1, 4])
Mean: tensor([[0.4815, 0.1062, 0.4204, 0.5198]], grad_fn=<MeanBackward1>)
"""
1.3 在torchrec中的嵌入表
- 在上边我们研究了torch中的嵌入表，接下来让我们研究TorchRec对嵌入表的改进，这样我们才能了解他这个框架是怎么加速的
演示代码：torchrec.EmbeddingBagConfig详解
核心参数详解
- 1、name (必填)
- 作用：嵌入表的名称，需唯一。
- 示例：“user_table” 或 “product_table”。
- 2、embedding_dim (必填)
- 作用：嵌入向量的维度（即每个ID映射到的稠密向量长度）。
- 示例：embedding_dim=64 表示每个ID对应一个64维向量。
- 3、num_embeddings (必填)
- 作用：嵌入表的容量（即最大ID值 + 1，如ID范围为0~4095，则设为4096）。
- 注意：需覆盖所有可能的输入ID，否则会导致越界错误。
- 4、feature_names (必填)
- 作用：关联多个输入特征到同一个嵌入表（如多个特征共享同一嵌入）
- 示例：feature_names=[“user_id”, “user_age_group”]。
- 5、pooling (必填)
- 作用：池化模式，定义如何将多个嵌入向量聚合成单个向量。
- 可选值：
- PoolingType.SUM：求和（默认，适合稀疏特征）。
- PoolingType.MEAN：取平均。
- PoolingType.NONE：不池化（需确保输入单个ID）。
- 6、optimizer (可选)
- 作用：指定优化器类型，用于更新该嵌入表的参数。
- 可选值：
- OptimizerType.SGD：随机梯度下降。
- OptimizerType.ADAM：Adam 优化器（默认）。
- OptimizerType.ADAGRAD：Adagrad。
- 7、lr (可选)
- 作用：优化器的学习率，默认继承全局学习率。
- 示例：lr=0.01。
- 8、weight_decay (可选)
- 作用：L2正则化权重衰减系数。
- 示例：weight_decay=1e-4。
- 9、enable_grad (可选)
- 作用：是否启用梯度计算（即是否参与反向传播）。
- 默认值：True。
- 用途：冻结部分嵌入表时可设为False。
- 10、sparse (可选)
- 作用：是否启用稀疏优化（如稀疏梯度更新）。
- 默认值：False。
- 适用场景：超大规模嵌入表（数十亿ID）时开启以提升性能。
- 11、table_size (可选)
- 作用：嵌入表分片数量（用于多GPU/多设备并行）。
- 示例：table_size=4 将表分到4个设备。
- 12、hash_size (可选)
- 作用：哈希嵌入的哈希表大小（用于动态处理未见过的ID）。
- 示例：hash_size=1000000 将未见ID哈希到100万空间。
# 我们先了解 torchrec.EmbeddingBagConfig
import torchrec
config = torchrec.EmbeddingBagConfig(
name="user_behavior",
embedding_dim=128,
num_embeddings=10_000_000, # 支持1000万用户
feature_names=["user_id", "behavior_sequence"],
pooling=torchrec.PoolingType.SUM,
optimizer=torchrec.OptimizerType.ADAM,
enable_grad=Ture,
lr=0.001,
weight_decay=0.01,
sparse=True, # 启用稀疏优化
table_size=4, # 分片到4个GPU
hash_size=20_000_000, # 哈希处理未见ID
dtype=torch.float16, # 正常是float32, 可以设置为float16显存占用减半
)
# 我们使用两个 embedding bags 构建了 EmbeddingBagCollection (EBC)
ebc = torchrec.EmbeddingBagCollection(
device="cpu",
tables=[
torchrec.EmbeddingBagConfig(
name="product_table",
embedding_dim=64,
num_embeddings=4096,
feature_names=["product"],
pooling=torchrec.PoolingType.SUM,
),
torchrec.EmbeddingBagConfig(
name="user_table",
embedding_dim=64,
num_embeddings=4096,
feature_names=["user"],
pooling=torchrec.PoolingType.SUM,
)
]
)
print(ebc.embedding_bags)
"""
打印结果如下：
ModuleDict(
(product_table): EmbeddingBag(4096, 64, mode='sum')
(user_table): EmbeddingBag(4096, 64, mode='sum')
)
"""
二、TorchRec中的输入输出数据类型
2.1 offsets的概念
演示代码：
# 以两个样本来举例
# 样本1 的长度为1。 样本2的长度为2
id_list_feature_lengths = torch.tensor([1, 2])
# 样本1的 值为5。 样本2 的值为 7和1
id_list_feature_values = torch.tensor([5, 7, 1])
# 计算每个样本的偏移量 cumsum函数就是按照行求和。 从 [1, 2]->[1,3]
id_list_feature_offsets = torch.cumsum(id_list_feature_lengths, dim=0)
print("Offsets: ", id_list_feature_offsets)
print("First Batch: ", id_list_feature_values[: id_list_feature_offsets[0]])
print(
"Second Batch: ",
id_list_feature_values[id_list_feature_offsets[0] : id_list_feature_offsets[1]],
)
"""
输出结果：
Offsets: tensor([1, 3])
First Batch: tensor([5])
Second Batch: tensor([7, 1])
"""
2.2 JaggedTensor
演示代码：
# JaggedTensor就是一个包含了上述操作的封装函数
from torchrec import JaggedTensor
jt = JaggedTensor(values=id_list_feature_values, lengths=id_list_feature_lengths)
# 可以自动的根据长度计算出偏移量，同理也可以根据偏移量计算出长度
print("Offsets: ", jt.offsets())
# 转化成稠密向量，也就是可以分开
print("List of Values: ", jt.to_dense())
print(jt)
"""
输出结果：
Offsets: tensor([0, 1, 3])
List of Values: [tensor([5]), tensor([7, 1])]
JaggedTensor({
[[5], [7, 1]]
})
"""
2.3 KeyedJaggedTensor
演示代码：
# JaggedTensor只能表示一个ids的一个特征，但是我们有多个特征在一个EmbeddingBagCollection中
from torchrec import KeyedJaggedTensor
# 我们有两个特征 product和user，我们看来构建JaggedTensor
product_jt = JaggedTensor(
values=torch.tensor([1, 2, 1, 5]), lengths=torch.tensor([3, 1])
)
user_jt = JaggedTensor(values=torch.tensor([2, 3, 4, 1]), lengths=torch.tensor([2, 2]))
kjt = KeyedJaggedTensor.from_jt_dict({"product": product_jt, "user": user_jt})
# 查看有多少个键
print("Keys: ", kjt.keys())
# 查看所有的长度
print("Lengths: ", kjt.lengths())
# 查看所有的值，会展成1D维度
print("Values: ", kjt.values())
# 转成原来的字典格式
print("to_dict: ", kjt.to_dict())
# 打印kjt
print(kjt)
"""
输出结果：
Keys: ['product', 'user']
Lengths: tensor([3, 1, 2, 2])
Values: tensor([1, 2, 1, 5, 2, 3, 4, 1])
to_dict: {'product': <torchrec.sparse.jagged_tensor.JaggedTensor object at 0x7953dd7428f0>, 'user': <torchrec.sparse.jagged_tensor.JaggedTensor object at 0x7953dd7434f0>}
KeyedJaggedTensor({
"product": [[1, 2, 1], [5]],
"user": [[2, 3], [4, 1]]
})
"""
2.4 运行结果
演示代码：
# 使用我们上边构建的嵌入表，可以自动进行嵌入，不用我们自己定义那么多的embedding层
result = ebc(kjt)
print(result)
# 结果是一个KeyedTensor，其中包含特征名称和嵌入结果的列表
print(result.keys())
# 形状是[2,128], 首先是batch_size是2, 然后我们有两个表“product”和“user”，每个表的维度都是64, 所以64+64就是128
print(result.values().shape)
# 用于确定属于每个特征的嵌入
result_dict = result.to_dict()
for key, embedding in result_dict.items():
print(key, embedding.shape)
"""
['product', 'user']
torch.Size([2, 128])
product torch.Size([2, 64])
user torch.Size([2, 64])
"""
总结
- 以上就是TorchRec中用到的嵌入方法以及一些函数的介绍。
扫一扫
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
