딥러닝 모델을 훈련시키고나면, 모델의 가중치를 저장하고 나중에 재사용할 수 있는 방법이 필요합니다. 파이썬의 PyTorch 라이브러리를 사용하면 모델의 가중치를 쉽게 저장하고 불러올 수 있습니다.

가중치 저장

PyTorch에서 모델의 가중치를 저장하기 위해서는 state_dict() 함수를 사용합니다. 다음은 저장하는 예제 코드입니다.

import torch

# 모델 정의
model = YourModel()

# 훈련 후 가중치 저장
torch.save(model.state_dict(), 'path/to/save/weights.pth')

위의 코드에서 YourModel()은 사용자가 정의한 모델입니다. torch.save() 함수를 이용하여 state_dict()를 모델의 가중치를 지정된 경로에 저장할 수 있습니다. 파일 경로는 사용자가 원하는 곳에 지정하면 됩니다.

가중치 불러오기

저장한 가중치를 다시 불러올 때는 load_state_dict() 함수를 사용합니다. 다음은 불러오는 예제 코드입니다.

import torch

# 모델 정의
model = YourModel()

# 불러올 가중치 파일 경로
weights_path = 'path/to/save/weights.pth'

# 가중치 불러오기
model.load_state_dict(torch.load(weights_path))

위의 코드에서 YourModel()은 사용자가 정의한 모델입니다. torch.load() 함수를 이용하여 저장된 가중치 파일을 불러온 뒤, load_state_dict()를 통해 모델의 가중치를 업데이트할 수 있습니다.

참고로, 가중치를 저장할 때와 불러올 때, 모델 클래스의 인스턴스를 생성해야 합니다. 이는 모델 구조를 미리 정의하고 해당 클래스를 인스턴스화하는 것을 의미합니다.

위의 방법을 사용하여 PyTorch에서 모델의 가중치를 저장하고 불러올 수 있습니다. 이를 통해 모델을 다른 환경이나 다른 프로젝트로 쉽게 이동할 수 있으며, 이전에 훈련한 모델을 재사용하여 시간과 노력을 절약할 수 있습니다.

참고 자료

• PyTorch 공식 문서 - Saving and Loading Models