bdusell / semiring-einsum

Generic PyTorch implementation of einsum that supports different semirings
https://bdusell.github.io/semiring-einsum/
MIT License
45 stars 7 forks source link

Question about how to choose block_size #1

Closed speedcell4 closed 3 years ago

speedcell4 commented 3 years ago

Hi, is there any guidance on how to choose block_size?

bdusell commented 3 years ago

The best way is to benchmark your code and find a setting that is reasonably fast without running out of memory. If you're running on GPU, I would recommend increasing block_size until you run out of memory. One thing to keep in mind is that block_size applies to all summed variables, so if you have two summed variables, the memory usage is proportional to block_size ** 2. It would be possible to use a separate block_size for each variable, but I have not implemented this (yet?).

speedcell4 commented 3 years ago

Thank you for your kind reply~