model.state_dict()でRuntimeError: CUDA error: out of memory が発生した時の解決方法

背景

　0番のGPUで学習し保存したモデルを1番のGPU上でロードしようとしたところ、以下のようにエラーが発生した。

import torch

model.state_dict(torch.load("model.pth")["state_dict"])
# RuntimeError: CUDA error: out of memory

エラーの原因

　どうやらtorch.load()はデフォルトで、モデルを保存した際にモデルが存在していたメモリ上で、モデルのロードを行う設定になっているらしい。自分の例で言うと、0番のGPU上のモデルを"model.pth"に保存したため、torch.load("model.pth")を行うと、0番のGPU上でモデルのロードが行なわれていた。
　0番のGPUは他のモデルの学習中でメモリが占有されていたため、モデルをロードすることができずに、out of memory のエラーを吐いたわけである。

解決方法

方法1

　torch.load()のパラメータmap_locationで、モデルをロードするdeviceを1番のGPUに指定することで解決した。

import torch

device = "cuda:1"
model.load_state_dict(torch.load("model.pth", map_location=device)["state_dict"])

　このようにして、モデルのロードを1番のGPU上で行うことが可能になった。

方法2

　この方法は、モデルをロードする時ではなく、モデルを保存する時の対策である。したがって、既に保存済みのモデルをロードしたい時には使えない。
　torch.save()でモデルを保存する前に、モデルを別のメモリに移す、という方法である。

import torch

model.to('cpu')
torch.save(model.state_dict(), 'model.pth')

　上の例では、モデルをCPUに移してから保存しているので、モデルをロードする際にはCPU上で読み込みが行われます。