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)가 제공되며 배치 크기에 따라 다중화 할 수 있습니다.