如何連接 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

好的，所以嘗試對這段代碼進行矢量化，到目前為止我已經想出了這個

# 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]]

我已盡力使代碼自我解釋，如果它不起作用，請告訴我

Answer 2

這是我使用 numpy 的高級索引的解決方案。 它的優點是直到最后一行才復制數據。

取陣

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 秒