CUDA не хватает памяти с умножением матрицы

Question

Я пытаюсь умножить 3 матрицы, но у меня заканчивается память CUDA.

# A: 3000 x 100 (~2MB)
# B: 100 x 100  (~0.05MB)
# C: 100 x 3MM  (~2GB)

A = np.random.randn(3000, 100)
B = np.random.randn(100, 100)
C = np.random.randn(100, 3e6)

A_gpu = torch.from_numpy(A).cuda()
B_gpu = torch.from_numpy(B).cuda()
C_gpu = torch.from_numpy(C).cuda()

R_gpu = (A_gpu @ B_gpu @ C_gpu)

Cuda запрашивает около 90 ГБ памяти для этой операции.Я не понимаю почему.

Answer 1

Умножая матрицы, ваш выходной размер будет 3000 x 3 000 000 матриц!поэтому, несмотря на то, что A и B относительно малы, вывод R ОГРОМНЫЙ: 9 G элементов.Более того, я подозреваю, что dtype из ваших матриц float64, а не float32 (потому что вы использовали numpy для их инициации).Следовательно, для каждого из элементов 9G R_gpu требуется 8 байтов;увеличивая объем памяти GPU не менее 72 ГБ только для R_gpu.Я подозреваю, что промежуточные результаты и некоторые другие вещи занимают немного больше памяти GPU.