如何将NumPy二维数组构成的数组转换为适配PyTorch的三维数组

问题场景 #

现有存储句子词嵌入结果的NumPy对象数组，由多个二维NumPy子数组组成，示例如下：

embedded_sentences = np.array([
   np.array([[1,2,3],[0,0,0]]),
   np.array([[10,20,30],[40,50,60]])
])

该数组初始shape为(2,)，对应2个句子，需要转换为shape=(2,2,3)的三维数值数组，目标结构如下：

np.array([
   [[1,2,3],[0,0,0]],
   [[10,20,30],[40,50,60]]
])

转换后的数组需要支持直接通过torch.from_numpy(embedded_sentences)转为PyTorch可用张量。
此前尝试np.vstack(embedded_sentences)、np.dstack(embedded_sentences)均无法得到预期结果；已验证可行的np.dstack+np.transpose组合写法不够简洁，需要更精简的实现方案。

最简实现方案 #

直接使用np.stack()即可一行完成转换，该函数默认沿新的第0轴堆叠输入数组，完全匹配需求：

embedded_sentences = np.stack(embedded_sentences)

转换后验证embedded_sentences.shape可得到输出(2,2,3)，结构和目标完全一致，无需额外转置操作，可直接传入PyTorch做张量转换。

不同API效果说明 #

• np.vstack()：沿现有第0轴拼接数组，会将输入的两个二维数组合并为shape=(4,3)的二维数组，不符合三维结构要求
• np.dstack()：沿第2轴拼接数组，输出为shape=(2,3,2)的数组，需要额外转置才能得到目标维度，步骤冗余
• np.stack()：会为输入的所有子数组新增外层堆叠维度，完整保留每个子数组的内部结构，是该场景下最适配的API。

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