Bagaimana cara mengetahui apakah data yang diteruskan ke GPU akan menyebabkan CUDA keluar dari memori atau tidak

Aug 17 2020

Saya menggunakan GPU untuk menjalankan beberapa model pembelajaran mendalam yang sangat besar, ketika saya memilih ukuran batch 8, itu bisa masuk ke dalam memori, tetapi jika saya menggunakan ukuran batch 16, itu akan menyebabkan kesalahan CUDA kehabisan memori, dan saya harus menghentikan prosesnya.

Pertanyaan saya adalah, sebelum benar-benar meneruskan data ke GPU, adakah cara agar saya dapat mengetahui seberapa besar data yang akan menempati GPU?

Misalnya, kode berikut adalah tentang bagaimana saya membuat dataloader pytorch dan meneruskan setiap batch dataloader ke GPU, dapatkah saya tahu seberapa besar itu sebelum saya memanggil 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)

Jawaban

3 David Aug 17 2020 at 15:36

Saya akan merekomendasikan menggunakan paket torchsummary di sini.

pip install torchsummary

dan sedang digunakan

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

Ini akan memberi Anda ukuran model, ukuran forward pass, dan ukuran backpass dalam MB untuk ukuran batch 1, dan Anda kemudian dapat mengalikannya dengan ukuran batch Anda.