[Pytorch] Pytorch 모델 불러오기(save, checkpoints, transfer learning)

모델을 학습시키고 난 후 학습 결과를 공유하고 싶을 때 모델을 저장하는 방식을 사용한다.

- 코랩을 사용해서 작업을 진행했을 때 코랩이 꺼지면 결과가 날아가는 대참사가 발생한다.

- 학습 결과를 저장 및 관리할 필요가 있다.

파이토치에서는 model.save() 함수를 이용해서 모델을 저장한다.

 

model.save()

- 학습의 결과를 저장하기 위한 함수이다.

- 모델 형태(아키텍처)와 파라미터를 저장하는 방식이 있다.

- 모델 학습 중간 과정의 저장을 통해 최선의 결과 모델을 선택한다.

- 모델을 공유할 수 있어 학습 재연성을 향상시킬 수 있다.

결론적으로 모델을 저장하는 과정을 통해 잃을 것은 없다는 소리..!

 

첫 번째로, 파라미터 저장 방식

- 모델의 파라미터만 저장할 수 있다.

- 모델의 파라미터를 표시하는 함수인 state_dict()를 활용해서 저장한다.

- 파라미터를 load 하는 함수인 load_state_dict()를 활용해서 불러온다.

 

코드로 봐보자!

먼저 모델을 작성하고

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=0),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))

        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=0),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=0),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))

        self.drop_out = nn.Dropout()
        self.fc1 = nn.Linear(3 * 3 * 64, 1000)
        self.fc2 = nn.Linear(1000, 1)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)

        out = out.view(out.size(0), -1)
        out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out
        
# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

- model.state_dict()로 모델의 파라미터를 확인할 수 있다.

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

Output:
Model's state_dict:
layer1.0.weight 	 torch.Size([16, 3, 3, 3])
layer1.0.bias 	 torch.Size([16])
layer1.1.weight 	 torch.Size([16])
layer1.1.bias 	 torch.Size([16])
layer1.1.running_mean 	 torch.Size([16])
layer1.1.running_var 	 torch.Size([16])
layer1.1.num_batches_tracked 	 torch.Size([])
layer2.0.weight 	 torch.Size([32, 16, 3, 3])
layer2.0.bias 	 torch.Size([32])
layer2.1.weight 	 torch.Size([32])
layer2.1.bias 	 torch.Size([32])
layer2.1.running_mean 	 torch.Size([32])
layer2.1.running_var 	 torch.Size([32])
layer2.1.num_batches_tracked 	 torch.Size([])
layer3.0.weight 	 torch.Size([64, 32, 3, 3])
layer3.0.bias 	 torch.Size([64])
layer3.1.weight 	 torch.Size([64])
layer3.1.bias 	 torch.Size([64])
layer3.1.running_mean 	 torch.Size([64])
layer3.1.running_var 	 torch.Size([64])
layer3.1.num_batches_tracked 	 torch.Size([])
fc1.weight 	 torch.Size([1000, 576])
fc1.bias 	 torch.Size([1000])
fc2.weight 	 torch.Size([1, 1000])
fc2.bias 	 torch.Size([1])

- model.state_dict()의 type은 collections.OrderDict이다.

 

모델의 파라미터를 저장하는 코드

MODEL_PATH ="saved"
if not os.path.exists(MODEL_PATH):
    os.makedirs(MODEL_PATH)
torch.save(model.state_dict(), 
           os.path.join(MODEL_PATH, "model.pt"))

- 모델의 파라미터를 불러올 때는 같은 모델의 형태에서만 load 할 수 있다.

new_model = TheModelClass()
new_model.load_state_dict(torch.load(os.path.join(
    MODEL_PATH, "model.pt")))

 

두 번째로 모델의 아키텍처와 함께 저장하는 방식

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

- MODEL_PATH에 "model_pickle.pt"로 저장된다.

 

 

모델의 아키텍처와 함께 불러오는 방식

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

Output:
TheModelClass(
  (layer1): Sequential(
    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(2, 2))
    (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU()
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (layer2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2))
    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU()
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (layer3): Sequential(
    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2))
    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU()
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (drop_out): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=576, out_features=1000, bias=True)
  (fc2): Linear(in_features=1000, out_features=1, bias=True)
)

- model.eval()을 수행하면 모델을 평가 모드로 설정하는 것이다.

 

-> 모델을 저장하는 방식으로 파라메터를 저장하는 방식을 많이 쓴다고 한다. 파이토치 문서에서도 이 방법을 좀 더 권장하는 것 같다. 상황에 맞게 쓰기!

 

파이토치는 keras에 비해 모델을 출력했을 때 깔끔하게 볼 수 없다. 

torchsummary 패키지를 활용해서 keras와 같은 형태로 모델을 출력할 수 있다. (파라미터의 개수, 레이어의 종류 넘버 등)

아래와 같이 사용한다.

!pip install torchsummary #torchsummary 설치

from torchsummary import summary
summary(model, (3, 224, 224)) #input_size 설정

Output:
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 16, 111, 111]             448
       BatchNorm2d-2         [-1, 16, 111, 111]              32
              ReLU-3         [-1, 16, 111, 111]               0
         MaxPool2d-4           [-1, 16, 55, 55]               0
            Conv2d-5           [-1, 32, 27, 27]           4,640
       BatchNorm2d-6           [-1, 32, 27, 27]              64
              ReLU-7           [-1, 32, 27, 27]               0
         MaxPool2d-8           [-1, 32, 13, 13]               0
            Conv2d-9             [-1, 64, 6, 6]          18,496
      BatchNorm2d-10             [-1, 64, 6, 6]             128
             ReLU-11             [-1, 64, 6, 6]               0
        MaxPool2d-12             [-1, 64, 3, 3]               0
          Dropout-13                  [-1, 576]               0
           Linear-14                 [-1, 1000]         577,000
           Linear-15                    [-1, 1]           1,001
================================================================
Total params: 601,809
Trainable params: 601,809
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 5.53
Params size (MB): 2.30
Estimated Total Size (MB): 8.40
----------------------------------------------------------------

- 위와 같이 깔끔하게 결과를 볼 수 있다.

- 모델을 gpu, cpu로 부르냐에 따라 달라지는 것 같기 때문에 에러가 뜬다면 확인해보자.

 

지금까지 학습 결과 저장 방법을 봤고,

학습의 중간 결과를 저장해서 최선의 결과를 선택하는 checkponts를 알아보자

 

checkpoints

- 학습의 중간 결과를 저장하여 최선의 결과를 선택한다.

- earlystopping 기법 사용 시 이전 학습의 결과물을 저장한다.

- loss와 metric 값을 지속적으로 확인하고 저장한다.

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

- colab에서 지속적인 학습을 위해 필요하다.

 

checkpoint 구현 코드

- torch.save() 안에 dict 형태로 epoch와 함께 모델 정보를 저장한다.

- f-string을 활용해서 파일 명으로 성능을 확인할 수 있게 파일을 저장한다.

for e in range(1, EPOCHS+1):
    epoch_loss = 0
    epoch_acc = 0
    for X_batch, y_batch in dataloader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device).type(torch.cuda.FloatTensor)
        
        optimizer.zero_grad()        
        y_pred = model(X_batch)
               
        loss = criterion(y_pred, y_batch.unsqueeze(1))
        acc = binary_acc(y_pred, y_batch.unsqueeze(1))
        
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
        epoch_acc += acc.item()
        
        
    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")
        

    print(f'Epoch {e+0:03}: | Loss: {epoch_loss/len(dataloader):.5f} | Acc: {epoch_acc/len(dataloader):.3f}')

checkpoint 불러오는 코드

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

- checkpoint 저장 방식은 자유롭게 할 수 있다.

 

마지막으로 남이 만든 모델을 쓰는 방법을 알아보자

 

Transfer learning

- 다른 데이터셋으로 만든 모델을 현재 데이터에 적용하는 방식이다.

- 일반적으로 대용량 데이터셋으로 만들어진 모델의 성능이 높다.

- 현재의 DL에서 가장 일반적인 학습 기법이다. (대부분)

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

파이토치의 TorchVison은 다양한 Pre-trained 모델을 제공해준다.

 

vgg19 모델 마지막에 linear layer를 추가해보자.

from torch import nn
from torchvision import models

class MyNewNet(nn.Module):   
    def __init__(self):
        super(MyNewNet, self).__init__()
        self.vgg19 = models.vgg19(pretrained=True)
        self.linear_layers = nn.Linear(1000, 1) #Linearlayer 추가


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

- pretrained 하려면 (pretrained=True)를 작성해야 한다.

 

- 추가 말고 기존의 layer를 변경하는 방식도 있다.

vgg.classifier._modules['6'] = torch.nn.Linear(4096, 1)

 

Freezing

- pretrained model을 활용 시 모델의 일부분을 frozen 시킨다.

- frozen 시키는 것은 파라미터의 값을 변경하지 않는 것이다.

#frozen
for param in my_model.parameters():
    param.requires_grad = False 

#Linear layer은 frozen 해제
for param in my_model.linear_layers.parameters():
    param.requires_grad = True

- linear_layers 파라미터만 변경된다.

 

 

모델을 잘 학습시키는 과정도 중요하지만 모델을 어떻게 저장하고 관리하는 것이 좋을지에 대해 고민하는 과정도 중요하다고 느꼈다.

그리고 transfer learning을 잘 활용하려면 기초를 잘 다지는 것이 중요할 것 같다고 느꼈다. 기초를 잘하면 응용도 잘할 테니까..!!!!! transfer learning은 시간 될 때 틈틈이 더 공부해야겠다. 뭔가 방법은 알았는데 어떻게 활용해야 할지 잘 모르겠는 느낌이다ㅠ.ㅠ