每次看到numpy数组和PyTorch Tensor的复杂的索引操作时就很懵,在这里整理一下numpy中的常用的索引方式(numpy中的索引方式和pytorch很相似),以例子为主。
整数索引
import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
b = a[1:3, 1:3]
c = a[1:3, [1, 2]]
d = a[[1, 2], [1, 2]]
e = a[[1, 2]]
f = a[1:3]
print('='*20, 'a', '='*20)
print(a)
print('='*20, 'b', '='*20)
print(b)
print('='*20, 'c', '='*20)
print(c)
print('='*20, 'd', '='*20)
print(d)
print('='*20, 'e', '='*20)
print(e)
print('='*20, 'f', '='*20)
print(f)
输出结果
==================== a ====================
[[1 2 3]
[4 5 6]
[7 8 9]]
==================== b ====================
[[5 6]
[8 9]]
==================== c ====================
[[5 6]
[8 9]]
==================== d ====================
[5 9]
==================== e ====================
[[4 5 6]
[7 8 9]]
==================== f ====================
[[4 5 6]
[7 8 9]]
bool 索引
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
b = a[a > 5]
print('='*20, 'a', '='*20)
print(a)
print('='*20, 'b', '='*20)
print(b)
输出结果
==================== a ====================
[[1 2 3]
[4 5 6]
[7 8 9]]
==================== b ====================
[6 7 8 9]