索引一个 3d-Torch.tensor，为每个样本沿特定轴的唯一索引答案

【问题标题】：Index a 3d-Torch.tensor with unique indexes for each sample along a specific axis索引一个 3d-Torch.tensor，为每个样本沿特定轴的唯一索引
【发布时间】：2022-09-22 21:11:03
【问题描述】：

我正在尝试根据使用 torch.nonzeroes （用于神经形态计算项目中的延迟解码器）获取的时间步长来索引 3 维火炬张量，但不确定如何进行。

我的输入张量 x 具有维度 [Timesteps, Batchsize, Classes] = [48, 256, 10]。

我正在使用 torch.nonzeroes 获取一个包含 256 个唯一时间步长的数组（批次中的每个样本 1 个），其中每个时间步长是该样本的第一个出现的非零值，沿着 48 个时间戳的时间轴（我意识到这是可能效率很低，但到目前为止还没有找到更好的方法）。

nonzeroes = torch.nonzero(x,as_tuple=True)
FirstSpike = []
for i in range(0,np.size(x.cpu().detach().numpy(),1)):
    nonzeroes = torch.nonzero(x[:,i,:],as_tuple=True)
    FirstSpike.append(nonzeroes[0][0].cpu().detach().numpy())

这将 FirstSpike 作为 256 个整数值的数组返回，用于每个时间步的索引/切片 ( [7, 9, 13, 43,...] )。

我想用它来得到一个尺寸为 [256, 10] 的数组 x2，其中 256 个批次的每个样本对应于其适当的时间片。（例如，样本 17 的时间步长为 7，而样本 57 的时间步长为 38）。我知道我可能可以在这样的 for 循环中获得它：

x2 = []
for i in range(0,np.size(x.cpu().detach().numpy(),1)):
    val = x[FirstSpike[i],i,:] #output dimension [1,10]
    x2.append(val) #Final x2 dimension [256,10]

但是，由于这是神经网络解码器的一部分，因此在记忆方面效率非常低。是否有更聪明的操作可以一次性做到这一点？

亲切的问候乔纳森

标签： numpy indexing pytorch numpy-ndarray numpy-slicing

【解决方案1】：

我用来获取第一个元素的一个技巧是像这样使用argmax

FirstSpike = torch.argmax(x != 0, axis=0)

【讨论】：