[PyTorch] AI 모델 저장 & 불러오기

PyTorch AI 모델 저장 & 불러오기에 대해 학습 후 정리한 포스팅 입니다.

 

이번 포스팅에서는 AI 모델 저장 및 불러오는 방법에 대해 학습합니다.

 

 

I) AI 모델을 저장하는 이유?

AI 모델을 저장하는 것은 재사용성, 배포, 버전 관리, 전이 학습 등의 이유로 중요합니다.

 

1. 재사용성: 훈련된 모델은 재사용될 수 있어 모델을 처음부터 다시 훈련할 필요성을 줄입니다

2. 배포: 모델을 훈련하고 저장한 후에는 다양한 환경에 배포하여 사용자에게 실시간 예측을 제공할 수 있습니다.
3. 버전 관리: 모델의 변화를 추적하고 이전 버전으로 되돌릴 수 있게 해줍니다.
4. 전이 학습: 저장된 모델은 다른 관련된 작업에서 사전 훈련된 모델을 시작점으로 전이 학습에 사용할 수 있습니다.

 

 

II) AI 모델 저장 (모델 전체 저장 VS 매개변수만 저장)

1) AI 모델 전체를 저장

• 모델 전체를 저장하면 파라미터 + optimizer, epoch 등 모든 상태를 저장한다.
• 공유, 협업 및 추론이나 추가 훈련을 위해 정확히 동일한 모델을 재생성해야 할 때 유용합니다.

import torch
import torchvision.models as models

# 모델의 모든 상태를 저장
torch.save(model, '../AI_model.pt')
torch.load('../AI_model.pt')

2) 매개변수만 저장

• 나중에 사용하거나 다른 프레임워크로 이동할 가중치만 저장하려는 경우에 유용합니다.
• 다른 가중치 값과 함께 아키텍처를 재사용하거나 세부 저장 공간이 제한적인 경우에 적합합니다.

import torch
import torchvision.models as models

# 모델의 매개변수만 저장
torch.save(model.state_dict(),'../AI_model_state.pt')
torch.load_stat_dict(torch.load('../AI_model_stat.pt')

 

 

 

III) AI 모델 불러오기

1) AI 모델 전체를 불러오기

import torch
import torchvision.models as models

# 전체 모델 불러오기
model = torch.load('saved_model.pt')

# 모델을 평가 모드로 설정
model.eval()

# 모델을 추론에 활용
outputs = model(inputs)

2) 매개변수만 불러오기

import torch
import torchvision.models as models

# 같은 모델 아키텍처의 새로운 인스턴스 생성
model = models.vgg16()

# 저장된 가중치를 모델에 불러오기
model.load_state_dict(torch.load('model_weights.pth'))

# 모델을 평가 모드로 설정
model.eval()

# 모델을 추론에 활용
outputs = model(inputs)

 

 

 

 

 

ref: 

[PyTorch] PyTorch 모델을 저장하는 방법 및 고려해야할 점

 

 

<Summary>

- PyTorch - AI 모델 저장 & 불러오기

 

 

 

*유의사항

- AI 산업기능요원이 공부하여 남긴 정리입니다.