对具有值 (1) 的索引求和

Question

我有一个形状为 (2, 2, 3) 的张量，例如：

a= tensor ([[[2, 0, 2],[1, 0, 0]],[[1, 0, 1],[0, 1, 0]]])

我想找到值 (1) 的索引，然后将 2 乘以这些索引的幂，最后将最后一维的结果相加，所以最终结果应该是这样的：

tensor ([[[],[2^0]],[[2^0+2^2],[2^1]]])

我的实际张量比这个例子大得多，所以我不想使用“for”循环，我必须使用广播......

我在想torch.pow(2,(a == 1).nonzero()).sum() 之类的东西，但它不起作用。我必须找到一种方法将(a== 1).nonzero() 仅应用于最后一个维度，有什么建议吗？谢谢。

Answer 1

将其更改为 NumPy 并应用：

a_n = a.numpy()
a_n = np.apply_along_axis(func1d=lambda x: np.sum(np.power(2,np.where(x==1))[0]), axis=2, arr=a_n)

a = torch.Tensor(a_n)

基本上，它在axis = 2中应用您想要的功能，我假设您的较大数组类似于(2,2,n)