Por que fico sem memória CUDA ao executar o modelo PyTorch [com memória GPU suficiente]?

Aug 17 2020

Estou fazendo esta pergunta porque estou treinando com sucesso uma rede de segmentação na minha GTX 2070 no laptop com VRAM de 8 GB e uso exatamente o mesmo código e exatamente as mesmas bibliotecas de software instaladas no meu PC desktop com uma GTX 1080TI e ainda assim funciona memória.

Por que isso acontece, considerando que:

  1. O mesmo Windows 10 + CUDA 10.1 + CUDNN 7.6.5.32 + Nvidia Driver 418.96 (vem junto com o CUDA 10.1) estão no laptop e no PC.

  2. O fato de o treinamento com TensorFlow 2.3 ser executado sem problemas na GPU do meu PC, mas falha na alocação de memória para treinamento apenas com PyTorch.

  3. PyTorch reconhece a GPU (imprime GTX 1080 TI) por meio do comando: print(torch.cuda.get_device_name(0))

  4. O PyTorch aloca memória ao executar este comando: torch.rand(20000, 20000).cuda() #allocated 1.5GB of VRAM.

Qual é a solução para isso?

Respostas

5 TimbusCalin Aug 17 2020 at 17:27

A maioria das pessoas (mesmo no tópico abaixo) pula para sugerir que diminuir o batch_size resolverá este problema. Na verdade, não neste caso. Por exemplo, seria ilógico para uma rede treinar em VRAM de 8 GB e ainda assim não conseguir treinar em VRAM de 11 GB, considerando que não havia outros aplicativos consumindo memória de vídeo no sistema com VRAM de 11 GB e exatamente a mesma configuração instalada e usava.

A razão pela qual isso aconteceu no meu caso foi que, ao usar o DataLoaderobjeto, defini um valor muito alto (12) para o workersparâmetro. Diminuir esse valor para 4 no meu caso resolveu o problema.

Na verdade, embora no final da discussão, a resposta fornecida por Yurasyk em https://github.com/pytorch/pytorch/issues/16417#issuecomment-599137646 apontou-me na direção certa.

Solução: Diminua o número de workersno PyTorch DataLoader. Embora eu não entenda exatamente por que essa solução funciona, presumo que ela esteja relacionada aos threads gerados nos bastidores para a busca de dados; pode ser que, em alguns processadores, esse erro apareça.