在Numpy中,我可以使用布尔和数值索引的任意组合来索引数组。
>>> x = np.arange(12).reshape(3, 4)
>>> x
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
>>> x[np.ix_([False, True, True], [0, 3])] = [[1, 2], [3, 4]]
array([[ 0, 1, 2, 3],
[ 1, 5, 6, 2],
[ 3, 9, 10, 4]])
假设x
和索引数组(布尔和数值)都是torchTensor(在GPU上,如果有区别的话)。是否有一种方法可以像Numpy解决方案中那样索引和变异x
?
2条答案
按热度按时间iyfjxgzm1#
np.ix_
所做的就是将布尔数组输入转换为整数索引数组(通过.nonzero
),然后引入单例维度以方便广播。例如:相当于
只要每个索引包含整数。
我不知道pytorch中有任何
ix_
等价物,但巧合的是,在用.nonzero
转换布尔索引后,得到的形状已经非常适合广播:然而,在更一般的上下文中,你可以像
np.ix_
一样,在索引Tensor中插入用于广播的单例维度。enxuqcxy2#
您可以使用
torch.meshgrid
: