如何在numpy中通过embedding将2D矩阵转换为3D矩阵？

核心思路 #

本质是按索引查表的向量化操作，不需要手动实现循环，主流数值计算框架都内置了对应优化算子，可直接调用得到目标形状的矩阵。

NumPy 实现 #

直接用高级索引或者np.take方法即可：

import numpy as np

# input_ids 形状为 [batch_size, seq_len]，embedding_matrix 形状为 [num_classes, embedding_size]
output = embedding_matrix[input_ids]
# 等价写法，大张量下性能更稳定
output = np.take(embedding_matrix, input_ids, axis=0)

输出output的形状自动为[batch_size, seq_len, embedding_size]，完全符合要求。

PyTorch 实现 #

可以用直接索引、index_select算子，或者直接调用预训练Embedding层：

import torch

# 方法1：直接索引，写法最简洁
output = embedding_matrix[input_ids]

# 方法2：index_select 对超大张量优化更好
output = torch.index_select(embedding_matrix, 0, input_ids.flatten()).reshape(input_ids.shape[0], input_ids.shape[1], -1)

# 方法3：模型内调用可以直接绑定预训练embedding权重
emb_layer = torch.nn.Embedding.from_pretrained(embedding_matrix, freeze=True)
output = emb_layer(input_ids)

TensorFlow 实现 #

调用tf.gather算子即可：

import tensorflow as tf

output = tf.gather(embedding_matrix, input_ids, axis=0)

输出形状自动匹配为[batch_size, seq_len, embedding_size]。

注意事项 #

• 要保证输入标签矩阵input_ids中的最大值不超过num_classes - 1，否则会触发索引越界错误
• 所有上述实现都是纯向量化操作，底层自动调用SIMD/GPU做并行加速，性能比Python层级的手动for循环高2~3个数量级

内容的提问来源于stack exchange，提问作者Kenenbek Arzymatov