如何连接numpy数组的动态切片？

Question

我有一个 TensorFlow 层，我用它来创建一个名为“cost_volume”的 4d 张量。

我的代码完成了工作。但是，由于嵌套的 for 循环而不是使用 numpy 内置函数，它相当慢。我需要它至少快 300 倍。如何将其转换为使用 numpy 内置函数且没有 for 循环的更高效的代码？

max_disparity = 10
layer = np.random.rand(2, 188, 621, 32)
cost_volume = np.random.rand(layer.shape[0], layer.shape[1], layer.shape[2],\
                                 max_disparity + 1, feature_size * 2)

for i in range(layer.shape[0]):
    for y in range(layer.shape[1]):
        for x in range(layer.shape[2]):
            for d in range(max_disparity + 1):
                if i == 0:
                    cost_volume[i][y][x][d] = np.concatenate((layer[0][y][x], \
                                                         layer[1][y][min(x + d, layer.shape[2] - 1)]))
                else:
                    cost_volume[i][y][x][d] = np.concatenate((layer[0][y][max(0, x - d)], \
                                                         layer[1][y][x]))

Answer 1

这是我使用 numpy 的 advanced indexing 的解决方案。它的优点是直到最后一行才复制数据。

取数组

max_disparity = 10
feature_size = 2
np.random.seed(123)
layer = np.random.rand(2, 188, 621, 32)

还有代码

I, Y, X, D, F = layer.shape[:-1] + (max_disparity + 1, feature_size)
left = ( ### tuple of index-arrrays for left side of concat
    np.zeros((I, Y, X, D), dtype=int), # index for axis 0
    np.array([[[[y for d in range(D)] for x in range(X)]
        for y in range(Y)] for i in range(I)]), # index for axis 1
    np.array([[[[x if i==0 else max(0, x-d) for d in range(D)] for x in range(X)]
        for y in range(Y)] for i in range(I)]), # index for axis 2
    )
right = ( ### tuple of index-arrrays for right side of concat
    np.ones((I, Y, X, D), dtype=int), # index for axis 0
    np.array([[[[y for d in range(D)] for x in range(X)]
        for y in range(Y)] for i in range(I)]), # index for axis 1
    np.array([[[[min(x+d, X-1) if i==0 else x for d in range(D)] for x in range(X)]
        for y in range(Y)] for i in range(I)]), # index for axis 2
    )
### now slice layer with both tuples and concatenate inner-most dim
cost_volume = np.concatenate([layer[left], layer[right]], axis=-1)

在配备 2.16 GHz 双核赛扬和 4 GB 的笔记本电脑上用时 41.2 秒

Answer 2

好的，所以尝试对这段代码进行矢量化，到目前为止我已经想出了这个

# just renaming some variables

Z = layer.shape[0]
Y = layer.shape[1]
X = layer.shape[2]
D = max_disparity + 1
F = feature_size

cost_volume = np.empty([Z, Y, X, D, F * 2])

实际代码：

cost_volume[0, ..., :F] = layer[0, :,  :, None]               # i == 0 and 1st half of concatenate
cost_volume[0, ..., F:] = layer[1, :, -1, None, None]         # i == 0 and 2nd half of concatenate, default min(x + d, X) == X 

cost_volume[1, ..., F:] = layer[1, :,  :, None]               # i == 1 and 2nd half of concatenate  
cost_volume[1, ..., :F] = layer[0, :,  0, None, None]         # i == 1 and 1st half of concatenate, default max(x - d, 0) == 0

Xi, Di = np.ogrid[0:X, 0:D]        # indices used in 2nd and 3rd axis 

m1 = Xi + Di < X                   # positions where min(x + d, X) == x + d  ie where default doesnt apply 
m2 = Xi - Di > 0                   # positions where max(x - d, 0) == x - d  ie where default doesnt apply 

cost_volume[0, :, m1, F:] = layer[1, :, (Xi + Di)[m1]]   # updating cost values where defaults should not apply
cost_volume[1, :, m2, :F] = layer[0, :, (Xi - Di)[m2]]

我已尽力使代码自我解释，如果它不起作用，请告诉我