sapo
(sapo)
November 30, 2019, 6:48pm
1
Hello,
I need to create a tensor where entry i,j
has value abs(i - j)
.
For now, the only solution I’ve found is by cycling each entry with for
loop:
A = torch.zeros(1000, 2000)
for i in range(1000):
for j in range(2000):
A[i, j] = abs(i - j)
Unfortunately this solution is extremely slow.
I need this tensor to use it as mask of another tensor.
Someone has any alternative idea?
albanD
(Alban D)
December 1, 2019, 11:43pm
2
I’m not sure about the exact code but you can use range and expand to have matrices like A[i, j] = i, forall j: range(1000).unsqueeze(1).expand(1000, 2000). And the same for the other one with unsqueeze(0)
.
Then you can substract them element-wise and take the absolute value of the result.
sapo
(sapo)
December 2, 2019, 10:27am
3
Really thanks, I was using python list multiplication with torch.stack
, but this is much faster.
On CPU:
In [22]: %timeit torch.stack([torch.arange(1000)] * 2000, 1)
42.3 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [23]: %timeit torch.arange(1000).unsqueeze(1).expand(1000, 2000)
13.1 µs ± 21.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
On GPU:
In [27]: %timeit torch.stack([torch.arange(1000).to(device)] * 2000, 1).to(device)
2.15 ms ± 6.66 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [30]: %timeit torch.arange(1000).to(device).unsqueeze(1).expand(1000, 2000).to(device)
43.8 µs ± 152 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
1 Like