GPUに渡されたデータがCUDAのメモリ不足を引き起こすかどうかを知る方法

Question

GPUを使用していくつかの非常に大きな深層学習モデルを実行しています。8のバッチサイズを選択すると、メモリに収まりますが、16のバッチサイズを使用すると、CUDAのメモリ不足エラーが発生します。そして私はプロセスを殺さなければなりません。

私の質問は、実際にデータをGPUに渡す前に、データがGPUで占める大きさを知る方法はありますか？

たとえば、次のコードは、pytorchデータローダーを作成してデータローダーの各バッチをGPUに渡す方法に関するものですが、呼び出す前にその大きさを知ることができますか？ batch.to(device)

train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
for step, batch in enumerate(train_dataloader):
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)
    b_labels = batch[2].to(device)

Answer 1

ここでtorchsummaryパッケージを使用することをお勧めします。

pip install torchsummary

および使用中

from torchsummary import summary
myModel.cuda()
summary(myModel, (shapeOfInput)) # where shapeOfInput is a tuple of the sample's dimensions

これにより、バッチサイズが1の場合、モデルのサイズ、フォワードパスのサイズ、およびバックパスのサイズがMB単位で表示され、バッチサイズで乗算できます。