在 PyTorch 中进行多个嵌入的更快方法？

Question

我正在开发一个基于火炬的库，用于使用表格数据集构建自动编码器。

一大特点是学习分类特征的嵌入。

然而，在实践中，同时训练许多嵌入层会造成一些减速。我正在使用 for 循环来执行此操作，并且在每次迭代中运行 for 循环是（我认为）导致速度变慢的原因。

在构建模型时，我将嵌入层与用户数据集中的每个分类特征相关联：

        for ft in self.categorical_fts:
            feature = self.categorical_fts[ft]
            n_cats = len(feature['cats']) + 1
            embed_dim = compute_embedding_size(n_cats)
            embed_layer = torch.nn.Embedding(n_cats, embed_dim)
            feature['embedding'] = embed_layer

然后，调用 .forward()：

        embeddings = []
        for i, ft in enumerate(self.categorical_fts):
            feature = self.categorical_fts[ft]
            emb = feature['embedding'](codes[i])
            embeddings.append(emb)

        #num and bin are numeric and binary features
        x = torch.cat(num + bin + embeddings, dim=1)

然后x进入密集层。

这完成了工作，但在每次前向传递期间运行此 for 循环确实会减慢训练速度，尤其是当数据集具有数十或数百个分类列时。

有人知道像这样矢量化的方法吗？谢谢！

更新：为了更清晰，我绘制了我如何将分类特征输入网络的草图。您可以看到每个分类列都有自己的嵌入矩阵，而数字特征在传递到前馈网络之前直接连接到它们的输出。

我们可以在不迭代每个嵌入矩阵的情况下做到这一点吗？

Answer 1

只需使用简单的索引[，尽管我不确定它是否足够快

这是所有特征具有相同 vocab_size 和嵌入暗淡的简化版本，但它应该适用于异构类别特征的情况

xdim = 240
embed_dim = 8
vocab_size = 64
embedding_table = torch.randn(size=(xdim, vocab_size, embed_dim))

batch_size = 32
x = torch.randint(vocab_size, size=(batch_size, xdim))

out = embedding_table[torch.arange(xdim), x]
out.shape  # (bz, xdim, embed_dim)

# unit test
i = np.random.randint(batch_size)
j = np.random.randint(xdim)

x_index = x[i][j]
w = embedding_table[j]

torch.allclose(w[x_index], out[i, j])