[PyTorch] 20. 모델 저장 및 불러오기

▮ model.state_dict()

모델의 파라미터를 반환해준다.

type을 확인해보면 OrderedDict임을 알 수 있다.

 

type(model.state_dict())

# collections.OrderedDict

for param_tensor in model.state_dict():
	print(param_tensor, "\t", model.state_dict()[param_tensor].size())

▮ torchsummary

모델의 정보를 확인하는 좀 더 좋은 방법이다.

파라미터의 개수나 모델의 크기를 확인할 때 유용하다.

 

from torchsummary import summary

summary(model, (3, 224, 224))

 

 

모델 save 및 load

▮ 모델의 파라미터를 save

 

# initalize model
model = ModelClass().cuda()

...

torch.save(model.state_dict(), os.path.join(MODEL_PATH, "model.pt")

▮ 모델의 파라미터를 load

 

# 동일한 모델을 먼저 만들어야한다!
new_model =  ModelClass()

# 같은 모델의 archotecture에서 파라미터만 load
new_model.load_state_dict(torch.load(os.path.join(MODEL_PATH, "model.pt")))

 

▮ 모델의 architecture와 파라미터를 저장

 

torch.save(model, os.path.join(MODEL_PATH, "model_pickle.pt"))

▮ 모델의 architecture와 파라미터를 load

 

model = torch.load(os.path.join(MODEL_PATH, "model_pickle.pt"))

 

 

checkpoints

최선의 결과를 선택하기 위해 학습의 중간 결과를 저장할 필요성이 있다.

또한 학습 과정에서 한 번에 전체 에폭(Epoch)을 반복할 수 없거나 모종의 이유로 학습이 중단될 경우 일정 에폭(Epooch)마다 학습된 결과를 저장하여 나중에 이어서 학습할 수도 있다.

 

checkpoint도 torch.save를 활용해 저장하는데, 다양한 정보를 저장하기 위해 Dictionary 형식으로 값을 저장한다.

 

일반적으로 epoch, loss, metric을 함께 저장하여 확인한다.

 

torch.save({
        'epoch': e,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': epoch_loss,
        }, f"saved/checkpoint_model_{e}_{epoch_loss/len(dataloader)}_{epoch_acc/len(dataloader)}.pt")

 

 

체크포인트를 불러오는 방법은 아래와 같다.

device = "cuda" if torch.cuda.is_available() else "cpu"

model = CustomModel().to(device)
criterion = nn.MSELoss().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.0001)

checkpoint = torch.load('./checkpoint_model_1_36.782_50.033.pt')
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
checkpoint_epoch = checkpoint["epoch"]

 

해당 epoch 이후로 이어서 학습을 시키기 위해 epoch 값도 불러온 후 해당 값부터 반복문으로 학습을 진행한다.

for epoch in range(checkpoint_epoch + 1, 10000):
    cost = 0.0

    for x, y in train_dataloader:
        x = x.to(device)
        y = y.to(device)

        output = model(x)
        loss = criterion(output, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        cost += loss
        if (epoch + 1) % 1000 == 0:
            print(f"Epoch : {epoch+1:4d}, Model : {list(model.parameters())}, Cost : {cost:.3f}")

 

 

 

Transfer learning

일반적으로 대용량 데이터셋으로 만들어진 모델이 성능이 좋은 편이다.

따라서 이런 대용량의 다른 데이터셋으로 만든 모델(Pretrained Model)을 현재 데이터에 적용하는 방법이 요즘에는 일반적인 학습 방법이다.

backbone architecture가 잘 학습된 모델에서 일부분만 변경하여 학습을 수행한다.

 

Freezing

Pretrained Model을 활용 시 모델의 일부분을 frozen 시킨다.

https://towardsdatascience.com/transfer-learning-from-pre-trained-models-f2393f124751

 

Pretrained model load

vgg16모델을 vgg에 할당

import torch
from torchvision import models

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

vgg = models.vgg16(pretrained=True).to(device)

 

vgg16모델은 1000개의 이미지를 분류하는 task로 학습되었지만

우리는 cat과 dog을 분류하는 task를 할 것이 때문에 모델에 아래와 같이 Linear Layer를 추가해준다.

class MyNewNet(nn.Module):   
    def __init__(self):
        super(MyNewNet, self).__init__()
        self.vgg19 = models.vgg19(pretrained=True)
        self.linear_layers = nn.Linear(1000, 1) # 모델에 마지막 Linear Layer을 추가

    # Defining the forward pass    
    def forward(self, x):
        x = self.vgg19(x)        
        return self.linear_layers(x)

 

마지막 레이어를 제외하고 나머지 레이어는 Freezing 시킨다.

my_model = MyNewNet()
my_model = my_model.to(device)

for param in my_model.parameters():
    param.requires_grad = False
    
for param in my_model.linear_layers.parameters():
	param.requires_grad = True

 

학습 전과 후의 파라미터를 확인하면 linear_layer만 변한것을 확인할 수 있다.

 

학습 전

 

학습 후

 

 

---

부스트캠프 AI Tech 교육 자료를 참고하였습니다.