Pytorch : 'RuntimeError: CUDA out of memory.'

파이토치 GPU 메모리 부족시 해결 방안

 

1. 메모리 제거

torch.cuda.empty_cache()

GPU에서 사용하지 않는 cached memory를 제거

 

 

2. eval 단계에서 gradient 계산안하게 설정

with torch.no_grad():

eval 단계에서는 back propagation을 수행하지 않기 때문에 gradient 계산을 하지 않아도 된다

torch.no_grad()로 설정 가능하다

 

KoBERT NSMC finetuning 에서도 아래와 같이 test accuracy 계산 시 적용 가능하다

 

https://github.com/SKTBrain/KoBERT/blob/master/scripts/NSMC/naver_review_classifications_pytorch_kobert.ipynb

GitHub - SKTBrain/KoBERT: Korean BERT pre-trained cased (KoBERT)

for e in range(num_epochs):
    train_acc = 0.0
    test_acc = 0.0
    model.train()
    '''
    생략
    '''
    print("epoch {} train acc {}".format(e+1, train_acc / (batch_id+1)))
    model.eval()
    with torch.no_grad(): # <---- 추가
        for batch_id, (token_ids, valid_length, segment_ids, label) in tqdm(enumerate(test_dataloader), total=len(test_dataloader)):
            token_ids = token_ids.long().to(device)
            segment_ids = segment_ids.long().to(device)
            valid_length= valid_length
            label = label.long().to(device)
            out = model(token_ids, valid_length, segment_ids)
            test_acc += calc_accuracy(out, label)
    print("epoch {} test acc {}".format(e+1, test_acc / (batch_id+1)))

 

3. 배치 사이즈 줄이기

배치사이즈를 줄여서 메모리 부하를 줄일 수 있다.

단, 논문에서 제시하는 하이퍼 파라미터에서 batch_size를 줄이면 오버피팅이나 값이 튀는 등 정확도에 차이가 생길 수 있으니 유의해야 한다