How to reduce the memory consumption of CUDA init in PyTorch?

Asked Dec 19 '22 at 03:53

Active Dec 19 '22 at 04:57

Viewed 358 times

I am using PyTorch on NVIDIA Jetson TX2 (GPU and CPU have shared memory), and have only about 2 Gb of free memory.

Specifications:

PyTorch v1.10.0
JetPack 4.6.1
CUDA 10.2

Whenever I try to use GPU, "torch.cuda.init()" would consume about 2Gb of memory. If I use only the CPU, the memory overhead would be only 180 Mb of memory.

I have been searching the net and found that the reason is a load of kernels. I understand that a lot of kernels are used for optimal computing, however, I can not even use my GPU due to it. Also, I think that the overwhelming majority of kernels would be never used by me.

Could you please say how to reduce the memory overhead if I am willing to sacrifice the performance as long as I could use GPU?

Would it be possible to reduce it to the CUDA bare minimum? If it is possible, how to accomplish it?

I am trying to train the relatively small model with Conv2D, Pooling, and Dense layers.

At least how to reduce the memory overhead, so I could train simple models with only Dense layers?

I also can not switch to other libraries due to other reasons.

edited Dec 19 '22 at 04:57

talonmies

70,661
34
192
269

asked Dec 19 '22 at 03:53

Mars

1

The docs for `torch.cuda.init()` says "Ordinary users should not need this, as all of PyTorch’s CUDA methods automatically initialize CUDA state on-demand." Have you tried not calling that function and just train your small model? – kmkurn Dec 19 '22 at 08:27
Yes, but it still takes 2 Gb – Mars Dec 20 '22 at 02:44
What is your batch size? Have you tried reducing it? – kmkurn Dec 20 '22 at 02:47
I used torch.rand(3,3).to("cuda") – Mars Dec 20 '22 at 04:43
So a single 3x3 float tensor takes 2GB GPU memory? That seems excessive. How did you get that 2GB? – kmkurn Dec 20 '22 at 05:07
If I initialize in CPU, it will take only 180 Mb of space. – Mars Dec 20 '22 at 13:13

How to reduce the memory consumption of CUDA init in PyTorch?

0 Answers0