从多维数组中获取第n项（python）

Question

我有一个形状为

(12, 19601, 1000)

的张量，我想只取第二个条目的每个第 n 个元素，我想要一个大小为

(12, 19601 / n, 1000)

的张量。有没有有效的方法来做到这一点？最终张量会更大，所以越快越好。预先感谢

Answer 1

为什么不使用切片作为索引？

import torch
x = torch.ones((12, 1024, 10))
x[:, ::2, :].shape  # torch.Size([12, 512, 10])
x[:, ::4, :].shape  # torch.Size([12, 256, 10])