[英]Speed up numpy integer-array indexing for depth
假设我有一个数组
[[0 2 1]
[1 0 1]
[2 1 1]]
我想将它转换为形式的张量
[[[1 0 0]
[0 1 0]
[0 0 0]]
[[0 0 1]
[1 0 1]
[0 1 1]]
[[0 1 0]
[0 0 0]
[1 0 0]]]
其中每个深度层(索引i
)是二进制掩码,显示i
在输入中出现的位置。
我已经为此编写了正确的代码,但对于任何使用都太慢了。 我可以用另一个向量化操作替换此函数中的循环吗?
def im2segmap(im, depth):
tensor = np.zeros((im.shape[0], im.shape[1], num_classes))
for c in range(depth):
rows, cols = np.argwhere(im==c).T
tensor[c, rows, cols] = 1
return tensor
使用broadcasting
-
(a==np.arange(num_classes)[:,None,None]).astype(int)
或者与builtin
外部比较 -
(np.equal.outer(range(num_classes),a)).astype(int)
如果你必须使用int
dtype或者通过完全跳过int
转换来保持为boolean
而使用uint8
以进一步提升。
样品运行 -
In [42]: a = np.array([[0,2,1],[1,0,1],[2,1,1]])
In [43]: num_classes = 3 # or depth
In [44]: (a==np.arange(num_classes)[:,None,None]).astype(int)
Out[44]:
array([[[1, 0, 0],
[0, 1, 0],
[0, 0, 0]],
[[0, 0, 1],
[1, 0, 1],
[0, 1, 1]],
[[0, 1, 0],
[0, 0, 0],
[1, 0, 0]]])
要将depth/num_classes
作为第三个dim,请扩展输入数组,然后与范围数组进行比较 -
(a[...,None]==np.arange(num_classes)).astype(int)
(np.equal.outer(im, range(num_classes))).astype(int)
(np.equal.outer(im, range(num_classes))).astype(np.uint8) # lower prec
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.