Làm thế nào để biết liệu dữ liệu được truyền đến GPU có gây ra CUDA hết bộ nhớ hay không

Question

Tôi đang sử dụng GPU để chạy một số mô hình học sâu rất lớn, khi tôi chọn kích thước lô là 8, nó có thể vừa với bộ nhớ, nhưng nếu tôi sử dụng kích thước lô là 16, nó sẽ gây ra lỗi hết bộ nhớ CUDA, và tôi phải giết quá trình.

Câu hỏi của tôi là, trước khi thực sự truyền dữ liệu vào GPU, có cách nào để tôi có thể biết dữ liệu sẽ chiếm dung lượng lớn như thế nào trong GPU không?

Ví dụ: đoạn mã sau nói về cách tôi tạo một bộ dữ liệu pytorch và chuyển từng lô bộ dữ liệu đến GPU, tôi có thể biết nó lớn như thế nào trước khi tôi gọi không 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

Tôi khuyên bạn nên sử dụng gói torchsummary ở đây.

pip install torchsummary

và đang sử dụng

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

Điều này sẽ cung cấp cho bạn kích thước của mô hình, kích thước của chuyển tiếp và kích thước của đường dẫn ngược tính bằng MB đối với kích thước lô là 1 và sau đó bạn có thể nhân ra theo kích thước lô của mình.