Какой самый быстрый способ скопировать значения из одного тензора в другой в PyTorch? - PullRequest
Новая
       43

Какой самый быстрый способ скопировать значения из одного тензора в другой в PyTorch?

4 голосов
/ 13 апреля 2019

Я экспериментирую с дилатацией в свертке, где я пытаюсь скопировать данные из одного 2D-тензора в другой 2D-тензор с помощью PyTorch.Я копирую значения из тензорного A в тензорный B таким образом, что каждый элемент A, который копируется в B, окружен n нулями.

Я уже пробовал использовать вложенные циклы for, что очень наивно.Производительность, очевидно, довольно плохая, когда я использую большое количество изображений в градациях серого в качестве входных данных.

for i in range(A.shape[0]):
   for j in range(A.shape[1]):
      B[n+i][n+j] = A[i][j]

Есть ли что-нибудь быстрее, что не требует использования циклов?

1 Ответ

5 голосов
/ 14 апреля 2019

Если я правильно понимаю ваш вопрос, вот более быстрая альтернатива, без каких-либо петель: 

# sample `n`
In [108]: n = 2

# sample tensor to work with
In [102]: A = torch.arange(start=1, end=5*4 + 1).view(5, -1)

In [103]: A
Out[103]: 
tensor([[ 1,  2,  3,  4],
        [ 5,  6,  7,  8],
        [ 9, 10, 11, 12],
        [13, 14, 15, 16],
        [17, 18, 19, 20]])

# our target tensor where we will copy values
# we need to multiply `n` by 2 since there are two axes
In [104]: B = torch.zeros(A.shape[0] + 2*n, A.shape[1] + 2*n)

# copy the values, at the center of the grid
# leaving `n` positions on the surrounding
In [106]: B[n:-n, n:-n] = A

# check whether we did it correctly
In [107]: B
Out[107]: 
tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  1.,  2.,  3.,  4.,  0.,  0.],
        [ 0.,  0.,  5.,  6.,  7.,  8.,  0.,  0.],
        [ 0.,  0.,  9., 10., 11., 12.,  0.,  0.],
        [ 0.,  0., 13., 14., 15., 16.,  0.,  0.],
        [ 0.,  0., 17., 18., 19., 20.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

Другой случай, когда n=3

In [118]: n = 3

# we need to multiply `n` by 2 since there are two axes
In [119]: B = torch.zeros(A.shape[0] + 2*n, A.shape[1] + 2*n)

# copy the values, at the center of the grid
# leaving `n` positions on the surrounding
In [120]: B[n:-n, n:-n] = A

In [121]: B
Out[121]: 
tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  1.,  2.,  3.,  4.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  5.,  6.,  7.,  8.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  9., 10., 11., 12.,  0.,  0.,  0.],
        [ 0.,  0.,  0., 13., 14., 15., 16.,  0.,  0.,  0.],
        [ 0.,  0.,  0., 17., 18., 19., 20.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

проверка работоспособности с вашим решением на основе loop: 

In [122]: n = 2
In [123]: B = torch.zeros(A.shape[0] + 2*n, A.shape[1] + 2*n)
In [124]: for i in range(A.shape[0]):
     ...:    for j in range(A.shape[1]):
     ...:       B[n+i][n+j] = A[i][j]
     ...:       

In [125]: B
Out[125]: 
tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  1.,  2.,  3.,  4.,  0.,  0.],
        [ 0.,  0.,  5.,  6.,  7.,  8.,  0.,  0.],
        [ 0.,  0.,  9., 10., 11., 12.,  0.,  0.],
        [ 0.,  0., 13., 14., 15., 16.,  0.,  0.],
        [ 0.,  0., 17., 18., 19., 20.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

время : 

# large sized input tensor
In [126]: A = torch.arange(start=1, end=5000*4 + 1).view(5000, -1)
In [127]: n = 2
In [132]: B = torch.zeros(A.shape[0] + 2*n, A.shape[1] + 2*n)
# loopy solution
In [133]: %%timeit
     ...: for i in range(A.shape[0]):
     ...:    for j in range(A.shape[1]):
     ...:       B[n+i][n+j] = A[i][j]
     ...:       
92.1 ms ± 434 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)


# clear out `B` again by reinitializing it.
In [128]: B = torch.zeros(A.shape[0] + 2*n, A.shape[1] + 2*n)

In [129]: %timeit B[n:-n, n:-n] = A
49.6 µs ± 239 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Из приведенных выше моментов мы можем видеть, что векторизованный подход на ~ 200x быстрее, чем решение на основе цикла.

...