HN Explorer

Ask HN: Why PyTorch einsum is significantly slower than transpose

8 points by rcshubhadeep ↗ HN
I have been tinkering with some DL models and wanted to implement part of it using PyTorch einsum. Before doing so I was wondering about the performance. I have been a bit skeptic as I believe there is a parsing (and even may be somewhat code generation) involved in implementation of einsum (I have never look under the hood of PyTorch or Numpy as to how is it implemented, so I may be completely wrong)

So to measure the performance, I created a simple benchmark of comparison. I created a Tensor with these dimensions (BATCH, X, Y). Like so -

a = torch.randn(10, 20, 30)

Then in Jupyter I did this

%%timeit

torch.einsum('b i j -> b j i', a)

AND

%%timeit

a.transpose(1, 2)

-----------------------

This is the result

5.43 µs ± 63.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each) [Einsum]

1.15 µs ± 2.51 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) [transpose]

Am I doing / reading something wrong? Is it a wrong way to benchmark? Or is it really true what I see, that einsum is order of magnitude slower than transpose?

4 comments

[ 0.23 ms ] story [ 17.6 ms ] thread
gus_massa
Have you tried with different sizes like

  a = torch.randn(10, 20, 30)
  a = torch.randn(20, 40, 60)
  a = torch.randn(30, 60, 90)
  ...
Is the "4µs" a constant difference or it's proportional to the size of the matrix?
rcshubhadeep
Tried a few sets, went until 2000, 3000 and the time gap between the two results remains more or less constant!... Not sure how to interpret it.
rcshubhadeep
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
My 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.)