จะรู้ได้อย่างไรว่าข้อมูลที่ส่งผ่านไปยัง GPU จะทำให้หน่วยความจำ CUDA ไม่เพียงพอหรือไม่

Question

ฉันใช้ GPU เพื่อเรียกใช้โมเดลการเรียนรู้เชิงลึกที่มีขนาดใหญ่มากเมื่อฉันเลือกขนาดแบทช์เป็น 8 มันสามารถใส่ลงในหน่วยความจำได้ แต่ถ้าฉันใช้ขนาดแบทช์เป็น 16 จะทำให้เกิดข้อผิดพลาด CUDA หน่วยความจำไม่เพียงพอ และฉันต้องฆ่ากระบวนการ

คำถามของฉันคือก่อนที่จะส่งข้อมูลไปยัง GPU จริงมีวิธีใดบ้างที่ฉันจะรู้ได้ว่าข้อมูลใน GPU จะมีขนาดใหญ่เพียงใด?

ตัวอย่างเช่นรหัสต่อไปนี้เป็นข้อมูลเกี่ยวกับวิธีสร้าง pytorch dataloader และส่งชุดข้อมูลแต่ละชุดไปยัง 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 จากนั้นคุณสามารถคูณด้วยขนาดแบทช์ของคุณได้