Pytorchのtensorが占有しているGPUのメモリを開放する方法

結論

GPUに移した変数をdelした後、torch.cuda.empty_cache()を叩くと良い。

検証1:delの後torch.cuda.empty_cache()を叩きGPUのメモリを確認

Python 3.7.8 | packaged by conda-forge | (default, Jul 31 2020, 02:25:08) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'1.6.0'
>>> import GPUtil
>>> GPUtil.showUtilization()
| ID | GPU | MEM |
------------------
|  0 |  0% |  7% |
>>> a = torch.rand(1, 1000, 1000, 1000).to("cuda")
>>> GPUtil.showUtilization()
| ID | GPU | MEM |
------------------
|  0 |  0% | 26% |
>>> del a
>>> torch.cuda.empty_cache()
>>> GPUtil.showUtilization()
| ID | GPU | MEM |
------------------
|  0 |  0% |  7% |

なぜdelの後にtorch.cuda.empty_cache()を叩くのか

delだけではキャッシュが残ってしまうから。
そのため、torch.cuda.empty_cache()でキャッシュも消す必要がある。

検証2:検証１のコードの合間にキャッシュとキャッシングアロケータのメモリを可視化

Python 3.7.8 | packaged by conda-forge | (default, Jul 31 2020, 02:25:08) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'1.6.0'
>>> device = torch.device("cuda:0")
>>> torch.cuda.get_device_properties(device=device).total_memory
24032378880
>>> torch.cuda.memory_reserved(device=device) # キャッシュ
0
>>> torch.cuda.memory_allocated(device=device)# キャッシングアロケータ
0
>>> a = torch.rand(1, 1000, 1000, 1000).to("cuda:0")
>>> torch.cuda.memory_reserved(device=device) # 大規模なtensorを作ったため0ではなくなっている。
4001366016
>>> torch.cuda.memory_allocated(device=device)# 大規模なtensorを作ったため0ではなくなっている。
4000000000
>>> del a
>>> torch.cuda.memory_reserved(device=device) # del aをした後でもキャッシュは変わらず
4001366016
>>> torch.cuda.memory_allocated(device=device)# キャッシングアロケータのメモリの占有は0になる
0
>>> torch.cuda.empty_cache()
>>> torch.cuda.memory_reserved(device=device) # キャッシュが0になる
0
>>> torch.cuda.memory_allocated(device=device)# 前回が０だったので変わらず
0

参考

https://pytorch.org/docs/1.6.0/cuda.html?highlight=torch%20cuda%20empty_cache#torch.cuda.empty_cache
https://pytorch.org/docs/stable/cuda.html?highlight=torch%20cuda%20memory_reserved#torch.cuda.memory_reserved