본문 바로가기

Deep Learning/PyTorch

PyTorch 학습 모델 저장 및 불러오기

딥러닝 학습이 이루어진 모델의 파라미터를 저장하고, 다시 이를 불러와서 학습 모델에 입력하는 방법은
다음의 코드를 통해 이행할 수 있다. 다만 주의 해야할 점은, 학습 파라미터를 저장하고 불러올 때,
해당 변수가 GPU상에 존재하는지 아니면 CPU상에 존재하는지를 유념해두고 저장 및 불러오기 절차를 거쳐야한다.
딥러닝을 학습할때는 대부분 GPU가 있다는 전제하에서 진행되나, CPU로만 진행될 수도 있으므로
이부분을 꼭 기억하자!!!
(P.S. pickle_protocol=4 는 Python 버전 3.4 이후에서부터 구현된 기능으로 큰 객체를 효율적으로 저장하게해줌.)