Pytorch train 과정에서 발생한 OOM(out of memory)

오류 발생 재현을 위한 코드

from transformers import AutoModel

device = "cuda"
model = Automodel.from_pretrained("klue/roberta-large").to(device)
kfold = 5

for fold in enumerate(kfold):
	
    model = deepcopy.copy(model)
    for idx, batch in enumerate(train_data_loader):
    	# train step code...

결론

위와 같이 Training 할 경우 에러가 발생한다. 사유는 Vram으로 모델을 보낸 상태에서 지속적으로 copy를 하기 때문에, 모델 내에서 모델이 중첩되므로, fold가 길어질수록, Vram 양에 따라 램이 터질 가능성이 높다.

 

해결을 위해 했던 시도

매 반복문마다 카피했던 모델을 지우고, 메모리를 정리해봤으나 실패

import gc
del model
torch.cuda.empty_cache()
gc.collect()

 

해결 방법

deepcopy 타이밍을 시스템 메모리 안에서 카피 후, Vram으로 보내면 정상적인 학습 절차가 가능하다.

from transformers import AutoModel

device = "cuda"
model = Automodel.from_pretrained("klue/roberta-large")#.to(device)
kfold = 5

for fold in enumerate(kfold):
	# copy 후 model 사용
    model = deepcopy.copy(model).to(device)
    for idx, batch in enumerate(train_data_loader):
    	# train step code...