PyTorch: ขณะโหลดข้อมูลแบบแบตช์โดยใช้ Dataloader วิธีถ่ายโอนข้อมูลไปยัง GPU โดยอัตโนมัติ

หากเราใช้คลาส Dataset และ Dataloader รวมกัน (ดังที่แสดงด้านล่าง) ฉันจะต้องโหลดข้อมูลลงใน GPU อย่างชัดเจนโดยใช้ .to() หรือ .cuda() มีวิธีสั่งให้ dataloader ทำโดยอัตโนมัติ/โดยปริยายหรือไม่?

รหัสเพื่อทำความเข้าใจ/สร้างสถานการณ์จำลอง:

from torch.utils.data import Dataset, DataLoader
import numpy as np

class DemoData(Dataset):
    def __init__(self, limit):
        super(DemoData, self).__init__()
        self.data = np.arange(limit)

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        return (self.data[idx], self.data[idx]*100)

demo = DemoData(100)

loader = DataLoader(demo, batch_size=50, shuffle=True)

for i, (i1, i2) in enumerate(loader):
    print('Batch Index: {}'.format(i))
    print('Shape of data item 1: {}; shape of data item 2: {}'.format(i1.shape, i2.shape))
    # i1, i2 = i1.to('cuda:0'), i2.to('cuda:0')
    print('Device of data item 1: {}; device of data item 2: {}\n'.format(i1.device, i2.device))

ซึ่งจะแสดงผลดังนี้ หมายเหตุ - หากไม่มีคำแนะนำในการถ่ายโอนอุปกรณ์ที่ชัดเจน ข้อมูลจะถูกโหลดไปยัง CPU:

Batch Index: 0
Shape of data item 1: torch.Size([50]); shape of data item 2: torch.Size([50])
Device of data item 1: cpu; device of data item 2: cpu

Batch Index: 1
Shape of data item 1: torch.Size([50]); shape of data item 2: torch.Size([50])
Device of data item 1: cpu; device of data item 2: cpu

วิธีแก้ไขที่เป็นไปได้อยู่ที่ repo PyTorch GitHub นี้ ปัญหา(ยังคงเปิดอยู่ในเวลาที่โพสต์คำถามนี้) แต่ฉันไม่สามารถทำให้มันทำงานได้เมื่อตัวโหลดข้อมูลต้องส่งคืนรายการข้อมูลหลายรายการ!


person anurag    schedule 28.01.2021    source แหล่งที่มา
comment
ไม่มีใครเจอปัญหานี้เลยเหรอ?   -  person anurag    schedule 01.02.2021
comment
พารามิเตอร์ collate-fn จะไม่มีประโยชน์เมื่อชุดข้อมูลส่งคืนค่าทูเพิลในแต่ละชุด ดังนั้นยังคงมองหาวิธีแก้ปัญหาที่ดีกว่า!   -  person anurag    schedule 02.02.2021


คำตอบ (1)


คุณสามารถแก้ไข collate_fn เพื่อจัดการหลายรายการพร้อมกัน:

from torch.utils.data.dataloader import default_collate

device = torch.device('cuda:0')  # or whatever device/cpu you like

# the new collate function is quite generic
loader = DataLoader(demo, batch_size=50, shuffle=True, 
                    collate_fn=lambda x: tuple(x_.to(device) for x_ in default_collate(x)))

โปรดทราบว่าหากคุณต้องการให้มีผู้ปฏิบัติงานหลายคนสำหรับตัวโหลดข้อมูล คุณจะต้องเพิ่ม

torch.multiprocessing.set_start_method('spawn')

หลังจาก if __name__ == '__main__' ของคุณ (ดูปัญหานี้)

ต้องบอกว่าดูเหมือนว่าการใช้ pin_memory=True ใน DataLoader ของคุณจะมีประสิทธิภาพมากกว่ามาก คุณได้ลองใช้ตัวเลือกนี้แล้วหรือยัง
ดูการปักหมุดหน่วยความจำ สำหรับ ข้อมูลมากกว่านี้.


อัปเดต (8 กุมภาพันธ์ 2021)
โพสต์นี้ทำให้ฉันดูเวลาที่ใช้ในข้อมูลต่อโมเดลระหว่างการฝึกอบรม ฉันเปรียบเทียบสามทางเลือก:

  1. DataLoader ทำงานบน CPU และหลังจากดึงข้อมูลแบทช์แล้วเท่านั้น ข้อมูลจะถูกย้ายไปยัง GPU
  2. เช่นเดียวกับ (1) แต่มี pin_memory=True ใน DataLoader
  3. วิธีที่เสนอให้ใช้ collate_fn เพื่อย้ายข้อมูลไปยัง GPU

จากการทดลองที่มีข้อจำกัดของฉัน ดูเหมือนว่าตัวเลือกที่สองจะทำงานได้ดีที่สุด (แต่ไม่ได้มีส่วนต่างมากนัก)
ตัวเลือกที่สามจำเป็นต้องยุ่งยากเกี่ยวกับ start_method ของกระบวนการตัวโหลดข้อมูล และดูเหมือนว่าจะมีค่าใช้จ่ายที่จุดเริ่มต้นของ แต่ละยุค

person Shai    schedule 02.02.2021
comment
คำสั่ง spawn ใช้ไม่ได้สำหรับฉัน เป็นคำสั่งเฉพาะสำหรับเวอร์ชันใดเวอร์ชันหนึ่งหรือไม่ - person anurag; 08.02.2021
comment
@anurag ฉันไม่รู้ - person Shai; 08.02.2021