[pytorch] 모델 일부분만 저장하기/불러오기

PyTorch에서는 모델의 일부분만 저장하거나 불러올 수 있다. 이 방법은 모델의 특정 부분에 대한 접근 권한을 제어하고 모델의 일부분만을 관리하고자 할 때 유용하다. 

모델의 특정 부분을 저장하고 불러오기 위해 PyTorch에서는 state_dict의 특정 키를 사용할 수 있다.

(state_dict는 모델의 가중치와 편향 등을 포함하는 사전(dictionary)이다)

다음은 모델의 특정 부분을 저장하고 불러오는 간단한 예제이다.

import torch
import torch.nn as nn

# 예제 모델 정의
class ExampleModel(nn.Module):
    def __init__(self):
        super(ExampleModel, self).__init__()
        self.layer1 = nn.Linear(10, 5)
        self.layer2 = nn.Linear(5, 2)

# 모델 인스턴스 생성
model = ExampleModel()

# 모델의 state_dict 가져오기
state_dict = model.state_dict()

# state_dict의 키 출력
print("Keys in the state_dict:")
for key in state_dict:
    print(key)


# 원하는 부분을 저장
torch.save({'layer1_weights': model.layer1.weight,
            'layer1_bias': model.layer1.bias}, 'partial_model.pth')

# 저장된 부분을 불러오기
partial_state_dict = torch.load('partial_model.pth')
model.layer1.weight.data.copy_(partial_state_dict['layer1_weights'])
model.layer1.bias.data.copy_(partial_state_dict['layer1_bias'])

• model.layer1.weight.data.copy_(partial_state_dict['layer1_weights'])
  • 이 코드는 model의 첫 번째 레이어인 layer1의 가중치를 partial_state_dict에서 가져온 layer1_weights의 값으로 복사한다.
  • data는 텐서의 데이터를 직접적으로 접근하기 위한 속성이며, .copy_()는 값을 복사하는 메서드이다.
  • 따라서, layer1_weights의 값을 layer1의 가중치에 복사한다는 뜻이다.
• 이런 식으로 모델의 특정 부분에 대한 값을 복사함으로써, 전체 모델이 아닌 일부분만을 업데이트하거나 불러올 수 있다. 

 

모델이 복잡해서 state dict의 특정 키값만 저장하고 불러오고 싶을 땐 다음과 같이 코드를 작성할 수 있다.

import torch

# 모델 인스턴스 생성
model = ...

# 저장할 때 특정 키 필터링
filtered_state_dict = {k: v for k, v in model.state_dict().items() if k.startswith('backbone')}

# 필터링된 state_dict 저장
torch.save(filtered_state_dict, 'backbone_model.pth')

# 저장된 state_dict 불러오기
filtered_state_dict = torch.load('backbone_model.pth')

# 현재 모델의 state_dict 가져오기
current_state_dict = model.state_dict()

# 불러온 state_dict를 현재 모델에 적용 (일치하는 키만)
for key, value in filtered_state_dict.items():
    current_state_dict[key] = value

# 모델에 적용된 state_dict 설정
model.load_state_dict(current_state_dict)

• model의 state dict 키 값중 'backbone'으로 시작하는 웨이트만 저장
• 불러온 모델 웨이트의 일부를 현재 모델의 state dict의 일치하는 키값에 적용