动态索引/选择numpy数组的维度

Question

刚刚在CNN上工作并且卡在张量算法上。

我希望能够遍历维度的列表或元组，并从该维度中选择X的一系列元素（多维数组），同时保留其他维度。

x = np.random.random((10,3,32,32)) #some multi dimensional array
dims = [2,3] #aka the 32s

#for a dimension in dims
#I want the array of numbers from i:i+window in that dimension

#something like
arr1 = x.index(i:i+3,axis = dim[0]) 
#returns shape 10,3,3,32

arr2 = arr1.index(i:i+3,axis = dim[1]) 
#returns shape 10,3,3,3

Answer 1

np.take应该适合您（阅读其文档）

In [237]: x=np.ones((10,3,32,32),int)
In [238]: dims=[2,3]
In [239]: arr1=x.take(range(1,1+3), axis=dims[0])
In [240]: arr1.shape
Out[240]: (10, 3, 3, 32)
In [241]: arr2=x.take(range(1,1+3), axis=dims[1])
In [242]: arr2.shape
Out[242]: (10, 3, 32, 3)

Answer 2

您可以尝试使用

进行切片

arr1 = x[:,:,i:i+3,:]

和

arr2 = arr1[:,:,:,i:i+3]

然后

形状

>>> x[:,:,i:i+3,:].shape
(10, 3, 3, 32)