opt-einsum-torch
There have been many implementations of Einstein's summation. numpy's numpy.einsum is the least efficient one as it only runs in single thread on CPU. PyTorch's torch.einsum works for both CPU and CUDA tensors. However, since there is no virtual CUDA memory, torch.einsum will run out of CUDA memory for large tensors.
This code aims at implementing a memory-efficient einsum function using PyTorch as the backend. This code also uses the opt_einsum package to optimizes the contraction path to achieve the minimal FLOPS.
Usage
from opt_einsum_torch import EinsumPlanner
import torch
# Some huge tensors
arr1, arr2 = ..., ...
ee = EinsumPlanner(torch.device('cuda:0'), cuda_mem_limit=0.9)
result = ee.einsum('ijk,jkl->il', arr1, arr2)
The resulting tensor result will be a PyTorch CPU tensor. You could convert it into numpy array by simply calling result.numpy().
Future works
- Support multiple GPUs.
- Memory efficient einsum kernels.
- CUDA data transfer profilers.
76 Dec 14, 2022
456 Dec 12, 2022
1.4k Dec 25, 2022
12 May 12, 2022
2 Jan 4, 2022
80 Dec 8, 2022
40 Dec 14, 2022
31 Oct 13, 2022
140 Dec 19, 2022
549 Dec 28, 2022
31 Sep 07, 2022
2 Feb 15, 2022
21 Dec 22, 2022
300 Dec 24, 2022
16 Apr 15, 2022
0 Dec 14, 2021
1.5k Dec 31, 2022
3 Jan 30, 2022
116 Oct 21, 2022
4 Aug 22, 2022
1 Jan 18, 2022
8 Sep 22, 2022
9 Sep 23, 2022
49 Nov 05, 2022
19 Dec 11, 2022
60 Jan 04, 2023
8.3k Dec 31, 2022
1.7k Jan 06, 2023
156 Dec 19, 2022