我正在进行卷积扩张实验,在该过程中,我尝试使用PyTorch将数据从一个2D张量复制到另一个2D张量。我复制从张值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]
是否有不需要循环使用的更快的东西?
如果我正确理解了您的问题,这是一个更快的选择,没有任何循环:
# 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
比基于循环的解决方案要快。