pytorch 모델을 cpu에서 테스트 후 gpu에서 돌릴때 out of memory 😭

cpu에서 테스트 후에 gpu에서 파인튜닝시, gpu 메모리가 증가하여 oom으로 중단되는 경우가 발생하였다.

OOM 발생 경우

1. for 문에서 data 사용

for 문에서 data를 할당 받는 경우, 데이터를 받을 때, 이전에 사용한 메모리를 저장하기 때문에 순간적으로 2배가 되어 oom이 발생할 수 있다.

2. loss를 torch.tensor를 그대로 이용하여 사용.

아래와 같이 total_loss를 누적하는 경우 그대로 loss를 torch.tensor로 이용하는 경우.

logit, loss= model(input)

total_loss += loss

3. Test 또는 Validation 시 메모리 오류

: train 시 괜찮은데 test에 오류가 나는 경우, test 시에는 기울기 계산을 하지 않기 위한 선언이 필요.

train code:      Variable(im_data.cuda())
evaluate code:   Variable(im_data.cuda(), volatile=True)

나의 해결 방안

위에 설명한 2번의 경우로 loss를 축적할 때, loss.item() 사용.

logit, loss= model(input)

# .item()으로 loss 값을 float 값만 받음
total_loss += loss.item() 

Reference

https://gldmg.tistory.com/108