我需要將 Pytorch 張量(類似于 numpy 陣列)與零的行和列交替。像這樣:
Input => [[ 1,2,3],
[ 4,5,6],
[ 7,8,9]]
output => [[ 1,0,2,0,3],
[ 0,0,0,0,0],
[ 4,0,5,0,6],
[ 0,0,0,0,0],
[ 7,0,8,0,9]]
我在這個問題中使用公認的答案,提出以下建議
def insert_zeros(a, N=1):
# a : Input array
# N : number of zeros to be inserted between consecutive rows and cols
out = np.zeros( (N 1)*np.array(a.shape)-N,dtype=a.dtype)
out[::N 1,::N 1] = a
return out
答案完美無缺,除了我需要在許多陣列上執行多次并且它所花費的時間已成為瓶頸。大部分時間是步長切片。
對于它的價值,我使用它的矩陣是 4D,矩陣的示例大小是 32x18x16x16,我只在最后兩個維度中插入備用行/列。
所以我的問題是,是否有另一種具有相同功能但時間更短的實作?
uj5u.com熱心網友回復:
我對Pytorch不熟悉,但是為了加速您提供的代碼,我認為JAX庫將有很大幫助。因此,如果:
import numpy as np
import jax
import jax.numpy as jnp
from functools import partial
a = np.arange(10000).reshape(100, 100)
b = jnp.array(a)
@partial(jax.jit, static_argnums=1)
def new(a, N):
out = jnp.zeros( (N 1)*np.array(a.shape)-N,dtype=a.dtype)
out = out.at[::N 1,::N 1].set(a)
return out
將在 GPU 上提高運行時間約 10 倍。它取決于陣列大小和N(大小的增加,更好的性能)。您可以根據迄今為止提出的 4 個答案(JAX 擊敗其他人)在我的Colab 鏈接上查看基準。
我相信,如果您可以針對您的問題進行調整,那么jax可以成為您案例的最佳庫之一(這是可能的)。
uj5u.com熱心網友回復:
我找到了一些方法來實作這個結果,并且索引方法似乎始終是最快的。
不過,其他方法可能會有一些改進,因為我試圖將它們從一維推廣到二維和任意數量的前導維度,并且可能沒有以最好的方式做到這一點。
編輯:另一種使用 numpy 的方法,而不是更快。
性能測驗(CPU):
In [4]: N, C, H, W = 11, 5, 128, 128
...: x = torch.rand(N, C, H, W)
...: k = 3
...:
...: x1 = interleave_index(x, k)
...: x2 = interleave_view(x, k)
...: x3 = interleave_einops(x, k)
...: x4 = interleave_convtranspose(x, k)
...: x4 = interleave_numpy(x, k)
...:
...: assert torch.all(x1 == x2)
...: assert torch.all(x2 == x3)
...: assert torch.all(x3 == x4)
...: assert torch.all(x4 == x5)
...:
...: %timeit interleave_index(x, k)
...: %timeit interleave_view(x, k)
...: %timeit interleave_einops(x, k)
...: %timeit interleave_convtranspose(x, k)
...: %timeit interleave_numpy(x, k)
9.51 ms ± 2.21 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
12.6 ms ± 4.98 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
23.3 ms ± 4.19 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
62.5 ms ± 19.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
50.6 ms ± 809 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
性能測驗(GPU):
(numpy方法未測驗)
...: ...
...: x = torch.rand(N, C, H, W, device="cuda")
...: ...
260 μs ± 1.92 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
861 μs ± 6.77 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
912 μs ± 14.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
429 μs ± 5.08 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
實作:
import torch
import torch.nn.functional as F
import einops
def interleave_index(x, k):
*cdims, Hin, Win = x.shape
Hout = (k 1) * (Hin - 1) 1
Wout = (k 1) * (Win - 1) 1
out = x.new_zeros(*cdims, Hout, Wout)
out[..., :: k 1, :: k 1] = x
return out
def interleave_view(x, k):
"""
From
https://discuss.pytorch.org/t/how-to-interleave-two-tensors-along-certain-dimension/11332/4
"""
*cdims, Hin, Win = x.shape
Hout = (k 1) * (Hin - 1) 1
Wout = (k 1) * (Win - 1) 1
zeros = [torch.zeros_like(x)] * k
out = torch.stack([x, *zeros], dim=-1).view(*cdims, Hin, Wout k)[..., :-k]
zeros = [torch.zeros_like(out)] * k
out = torch.stack([out, *zeros], dim=-2).view(*cdims, Hout k, Wout)[..., :-k, :]
return out
def interleave_einops(x, k):
"""
From
https://discuss.pytorch.org/t/how-to-interleave-two-tensors-along-certain-dimension/11332/6
"""
zeros = [torch.zeros_like(x)] * k
out = einops.rearrange([x, *zeros], "t ... h w -> ... h (w t)")[..., :-k]
zeros = [torch.zeros_like(out)] * k
out = einops.rearrange([out, *zeros], "t ... h w -> ... (h t) w")[..., :-k, :]
return out
def interleave_convtranspose(x, k):
"""
From
https://github.com/pytorch/pytorch/issues/7911#issuecomment-515493009
"""
C = x.shape[-3]
weight=x.new_ones(C, 1, 1, 1)
return F.conv_transpose2d(x, weight=weight, stride=k 1, groups=C)
def interleave_numpy(x, k):
"""
From https://stackoverflow.com/a/53179919
"""
pos = np.repeat(np.arange(1, x.shape[-1]), k)
out = np.insert(x, pos, 0, axis=-1)
pos = np.repeat(np.arange(1, x.shape[-2]), k)
out = np.insert(out, pos, 0, axis=-2)
return out
uj5u.com熱心網友回復:
由于您事先知道陣列的大小,因此優化的第一步是out在函式外部創建陣列。然后,嘗試jit 編譯函式并在陣列numba上就地作業。這比您發布out的版本實作了 5 倍的加速。numpy
import numpy as np
from numba import njit
@njit
def insert_zeros_n(a, out, N=1):
for i in range(a.shape[0]):
for j in range(a.shape[1]):
out[2*i,2*j] = a[i,j]
N并使用指定的and呼叫它a:
N = 1
a = np.arange(16*16).reshape(16, 16)
out = np.zeros( (N 1)*np.array(a.shape)-N,dtype=a.dtype)
insert_zeros_n(a,out)
uj5u.com熱心網友回復:
封裝任何N,如何使用numpy.kron4D 輸入,
a = np.arange(1, 19).reshape((1, 2, 3, 3))
print(a)
# array([[[[ 1, 2, 3],
# [ 4, 5, 6],
# [ 7, 8, 9]],
#
# [[10, 11, 12],
# [13, 14, 15],
# [16, 17, 18]]]])
def interleave_kron(a, N=1):
n = N 1
return np.kron(
a, np.hstack((1, np.zeros(pow(n, 2) - 1))).reshape((1, 1, n, n))
)[..., :-N, :-N]
np.hstack((1, np.zeros(pow(n, 2) - 1))).reshape((1, 1, n, n))為了性能,可以將where一次全部外部化/默認。
進而
>>> interleave_kron(a, N=2)
array([[[[ 1., 0., 0., 2., 0., 0., 3.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 4., 0., 0., 5., 0., 0., 6.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 7., 0., 0., 8., 0., 0., 9.]],
[[10., 0., 0., 11., 0., 0., 12.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0.],
[13., 0., 0., 14., 0., 0., 15.],
[ 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0.],
[16., 0., 0., 17., 0., 0., 18.]]]])
?
轉載請註明出處,本文鏈接:https://www.uj5u.com/qukuanlian/479491.html
