按索引对 ndarray 中的元素进行分组

Question

我有一个包含 1000 张图像的图像数据集，我已经为其创建了嵌入。每个嵌入（每个图像有 512 个嵌入，具有 256 维向量）是一个形状为 (512, 256) 的 ndarray，因此总的数组形状为 (1000, 512, 256)。

现在，从每个图像 (1000) 中，我想为第一个嵌入创建一组观察，可用的 512 个，并从每个图像中收集这个嵌入。然后我想对第二个嵌入、第三个、第四个，直到第 512 个嵌入执行此操作。

我将如何着手创建这些组？

Answer 1

您可以按如下方式实现：

groups = []

for i in range(512):
    # Select the i-th embedding from each image
    group = embeddings[:, i, :]
    groups.append(group)

groups = np.array(groups)