CuPy running out of memory

Question

I have been testing the CuPy library and done a simple matrix multiplication using einsum:

C = cp.einsum('pqrs,rs->pq', A, B)

Dimensions of A and B are, (41, 41, 41, 41) (41, 41), receptively. I also checked their sizes, which are 22606088 bytes, 13448 bytes.

While running the code, I am getting the following error message:
OutOfMemoryError: out of memory to allocate 38000834048 bytes (total 38023468032 bytes)

It indicates that I am running out of memory. Is there any option to sent data partially to the device and perform operations in terms of batches?

Answer 1

I think there is no option to send data partially for one-array.

And I faced same issue before, this may be caused because the cupy einsum efficiency is not optimized yet. https://github.com/cupy/cupy/issues/19#issuecomment-322972682

If you can try replacing your einsum function by using transpose, reshape and matmul etc, please try those.

I guess

C = cp.einsum('pqrs,rs->pq', A, B)

is equivalent to

p, q, r, s = A.shape
A = cp.reshape(A, (p, q, r*s))
B = cp.reshape(B, (1, 1, r*s))
C = cp.sum(A * B, axis=2)