4 ms·
To give you an example a = torch.randn(100, 5000, 8000) %%timeit torch.einsum('b i j -> b j i', a) ---- > 5.3 µs ± 11 ns per loop (mean ± std. dev. of 7 runs
by rcshubhadeep 5y ago
To give you an example
a = torch.randn(100, 5000, 8000)
%%timeit
torch.einsum('b i j -> b j i', a) ---- > 5.3 µs ± 11 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
%%timeit
a.transpose(1, 2) ---- > 1.16 µs ± 3.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
Tensors were not put to GPU. These are CPU based result.
- gus_massa 5y agoMy guess is that 'b i j -> b j i' is parsed in each iteration, and that is slower that using the numbers of the index explicitly. It's very strange that the time remains 5µs and 1µs in spite of the increase of the size. 100 * 5000 * 8000 is big enough to be easy to spot in the times. Are you sure the compiler is doing the calculation? (For example in Racket sometimes the compiler notice that you are not going to use the result, and the calculation will not raise any error, so the compiler just removes it. So for the benchmarks it's necessary to add something that use the result.) EDIT: Is the library using a lazy representation of the random matrix? (This would be very strange. How much time does it take to generate the initial value?) Is the library using a lazy representation of the transposed matrix? (I've seen this before a few times, sometimes the transposed version share the memory with the original one.)