[Pytorch] 메모리 효율적으로 사용하기

딥러닝으로 task를 진행할 때 여러 모델을 연속적으로 사용해야 하는 경우가 종종 있습니다. 이 때 pytorch에서 메모리를 좀 더 효율적으로 사용할 수 있는 방법에 대해 작성해보겠습니다.

 

1) 먼저 사용할 모델을 gpu로 옮겨줍니다.

model.to($device)

 

2) 모델의 사용이 끝나면 바로 cpu로 옮겨줍니다.

model.to('cpu')

 

3) garbage collect를 실행하고, 메모리를 비워줍니다.

gc.collect()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()