pytorchでCUDAのMemoryAllocation時のフラグメンテーションを防ぐ

pytorchで実装したモデルを学習中、最後のepochの途中で以下のようなエラーが出た。

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 640.00 MiB. GPU 0 has a total capacity of 23.65 GiB of which 59.69 MiB is free. Process 1331364 has 23.59 GiB memory in use. Of the allocated memory 22.84 GiB is allocated by PyTorch, and 255.30 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

ここではPYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueにすることで回避できる可能性があることを示している。
[1]によれば、この設定はCUDAのメモリ割り当てのfragmentationを防ぐことができるとのこと。
通常、CUDAのメモリ割り当ては前回割り当てた領域を再利用し、足りない分は新たに要求する振る舞いをする。毎回のイテレーションでテンソルサイズが全く同じ場合は問題が生じないが、batchサイズが毎回異なる場合など（サンプルを条件でフィルタリングした場合など）では、前回要求したサイズよりも必要なCUDAメモリサイズが大きくなる場合がある。このとき、毎回のメモリの要求サイズはオリジナルの要求サイズの1倍、2倍、4倍...と増えていくが、その都度実際に使われていない隙間ができるので、メモリのフラグ面テーションが生じる。

一方で、expandable_segments:Trueに設定するとプロセスが最初に一回割り当てたのち、必要に応じて拡張する方式をとるようになるので、使われないメモリ領域の隙間がなくなり、メモリ割り当て違反が生じにくくなる。

実際に環境変数をexportして実験してみると、CUDAのメモリの使用率が23.0GB->21.6GBと1.4GB程度減少するのを確認した。

#!/usr/bin/bash
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python train.py

Reference

• [1] https://pytorch.org/docs/stable/notes/cuda.html#environment-variables