Пики памяти графического процессора Conv1d во время обратной опоры. на меньших входах ⇐ Python

[Гость]

(Краткий воспроизводимый код ниже.)

Очень странное поведение при использовании torch.Conv1d: при подаче меньшего входного сигнала (ниже некоторого порога) использование памяти графического процессора резко возрастает в обратном направлении — на порядок или больше.

Мы предполагаем, что это связано с тем, что torch/cuda использует разные алгоритмы свертки в зависимости от размеров и доступной памяти; проблема в том, что это приводит к очень нежелательным ошибкам OOM во время выполнения.

импортировать факел из импорта фонарика nn %load_ext pytorch_memlab base_ch = 512 d_pos = 64 защита print_gpu_mem_usage(prefix=""): print(f"{prefix}Пиковая память: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} ГБ" f" | {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} ГБ" f" (Текущий: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} ГБ" f" | {torch.cuda.memory_reserved() / 1024 ** 3:.2f} ГБ)") защита изолированный_конв (в): samp_conv = nn.Conv1d(base_ch + d_pos, 2 * base_ch, kernel_size=1, дополнение='valid').cuda() мн = samp_conv(v).mean() мин.назад() %mlrun -f изолированный_конв изолированный_конв(torch.rand(5000, base_ch+d_pos, 11).cuda())
:active_bytes: :reserved_bytes: :line: :код: все все ----пик---- -----пик----- ---- ---- 108,00 млн 108,00 6 def изолированный_конв(v): 328,00 млн 346,00 млн 7 mn = nn.Conv1d(.....) 542,00 млн 562,00 млн 8 mn.backward()
однако, если переключить n выборок с 5000 на 4000, ситуация взорвется:

%mlrun -f изолированный_конв изолированный_конв(torch.rand(4000, base_ch+d_pos, 11).cuda())
:active_bytes: :reserved_bytes: :line: :код: все все ----пик---- -----пик----- ---- ---- 86,00 млн 86.00 6 def изолированный_конв(v): 260,00 млн 280,00 млн 7 mn = nn.Conv1d(.....) 8.07G 8,25 ГБ 8 mn.backward()
То же самое произойдет, если я проверю их в противоположном порядке.

Это работает в Docker, поэтому, если вы не можете воспроизвести его в следующих версиях, я могу поделиться Dockerfile.

факел == 2.0.1 pytorch-memlab == 0.3.0 Нвидиа 2080Ти Версия драйвера: 525.105.17
Версия CUDA: 12.0