在PyTorch中向量化张量到切片

时间:2019-09-27 20:14:13

标签: python arrays vectorization pytorch tensor

我正在尝试向量化形式的切片分配

for i in range(a.shape[1]):
    for j in range(a.shape[2]):
        a[:,i,j,:,i:i+b.shape[2],j:j+b.shape[3]] = b

其中b本身是一个数组。这是因为嵌套的Python循环效率低下,并且占用了大部分运行时间。有办法吗?

对于更简单的情况,请考虑以下事项:

for i in range(a.shape[1]):
    a[:,i,:,i:i+b.shape[2]] = b

ba可能看起来像这样:

enter image description here

您可以看到所得矩阵的对角线“滑动”结构。

2 个答案:

答案 0 :(得分:1)

我们可以利用基于np.lib.stride_tricks.as_stridedscikit-image's view_as_windows来将滑动的窗口视图获取到输入的0s填充版本中,成为一个视图将在内存和性能上提高效率。 More info on use of as_strided based view_as_windows

因此,对于较简单的情况,应该是-

from skimage.util.shape import view_as_windows

def sliding_2D_windows(b, outshp_axis1):
    # outshp_axis1 is desired output's shape along axis=1
    n = outshp_axis1-1
    b1 = np.pad(b,((0,0),(0,0),(n,n)),'constant')
    w_shp = (1,b1.shape[1],b.shape[2]+n)
    return view_as_windows(b1,w_shp)[...,0,::-1,0,:,:]

样品运行-

In [192]: b
Out[192]: 
array([[[54, 57, 74, 77],
        [77, 19, 93, 31],
        [46, 97, 80, 98]],

       [[98, 22, 68, 75],
        [49, 97, 56, 98],
        [91, 47, 35, 87]]])

In [193]: sliding_2D_windows(b, outshp_axis1=3)
Out[193]: 
array([[[[54, 57, 74, 77,  0,  0],
         [77, 19, 93, 31,  0,  0],
         [46, 97, 80, 98,  0,  0]],

        [[ 0, 54, 57, 74, 77,  0],
         [ 0, 77, 19, 93, 31,  0],
         [ 0, 46, 97, 80, 98,  0]],

        [[ 0,  0, 54, 57, 74, 77],
         [ 0,  0, 77, 19, 93, 31],
         [ 0,  0, 46, 97, 80, 98]]],


       [[[98, 22, 68, 75,  0,  0],
         [49, 97, 56, 98,  0,  0],
         [91, 47, 35, 87,  0,  0]],

        ....

        [[ 0,  0, 98, 22, 68, 75],
         [ 0,  0, 49, 97, 56, 98],
         [ 0,  0, 91, 47, 35, 87]]]])

答案 1 :(得分:0)

假设b的形状为(2,3,x1),而a的形状为(2,x2-x1+1,3,x2)。在您的屏幕截图中,我们可以推断出x1=4x2=6

import numpy as np

b_shape = (2,3,4)
a_shape = (2,3,3,6)

b = np.arange(1,25).reshape(b_shape)
#array([[[ 1,  2,  3,  4],
#        [ 5,  6,  7,  8],
#        [ 9, 10, 11, 12]],
#
#       [[13, 14, 15, 16],
#        [17, 18, 19, 20],
#        [21, 22, 23, 24]]])

c = np.pad(b, (*[(0,0) for _ in range(len(b_shape[:-1]))], (0,a_shape[-1]-b_shape[-1])), 'constant')

#array([[[ 1,  2,  3,  4,  0,  0],
#        [ 5,  6,  7,  8,  0,  0],
#        [ 9, 10, 11, 12,  0,  0]],
#
#       [[13, 14, 15, 16,  0,  0],
#        [17, 18, 19, 20,  0,  0],
#        [21, 22, 23, 24,  0,  0]]])

a = np.stack([np.roll(c, shift=i) for i in range(a_shape[-1]-b_shape[-1]+1)], axis=1)

# array([[[[ 1,  2,  3,  4,  0,  0],
#          [ 5,  6,  7,  8,  0,  0],
#          [ 9, 10, 11, 12,  0,  0]],

#         [[ 0,  1,  2,  3,  4,  0],
#          [ 0,  5,  6,  7,  8,  0],
#          [ 0,  9, 10, 11, 12,  0]],

#         [[ 0,  0,  1,  2,  3,  4],
#          [ 0,  0,  5,  6,  7,  8],
#          [ 0,  0,  9, 10, 11, 12]]],


#        [[[13, 14, 15, 16,  0,  0],
#          [17, 18, 19, 20,  0,  0],
#          [21, 22, 23, 24,  0,  0]],

#         [[ 0, 13, 14, 15, 16,  0],
#          [ 0, 17, 18, 19, 20,  0],
#          [ 0, 21, 22, 23, 24,  0]],

#         [[ 0,  0, 13, 14, 15, 16],
#          [ 0,  0, 17, 18, 19, 20],
#          [ 0,  0, 21, 22, 23, 24]]]])