Pytorch 학습하기

김동환·2023년 3월 16일
0

AI_tech_5기

목록 보기
3/18
post-custom-banner

Model 정보

Model의 정보를 보기 위해서는 model의 state_dict를 보면 된다. 리턴은 order_dict 타입으로

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

이렇게 하면 다음과 같다.

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])

아니면 summary를 이용해도 된다. 다만 summary를 이용할 경우 input의 모양을 넣어줘야한다.

from torchsummary import summary
summary(model, (3, 224, 224))

Model 저장 및 불러오기

state dict

저장하기

torch.save(model.state_dict(), "model.pt"))

불러오기

new_model.load_state_dict(torch.load("model.pt")))

pickle

저장하기

torch.save(model,"model_pickle.pt"))

불러오기

model = torch.load("model_pickle.pt"))

Transfer Learning

  • 다른 데이터셋으로 만든 모델을 현재 데이터에 적용
  • backbone architecture가 잘 학습된 모델에서 일부만 변경

Freezing

pretrained model 활용 시 모델의 일부분은 frozen 시킴
torchvision에서 모델을 가져올 수 있는데 이번에는 vgg 모델을 가져와보겠다. pretrained 옵션을 줘야 w도 같이 온다.

model 수정

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

그리고 클래스에서 vgg19모델을 받고 forward에서 새로 추가한 linear 레이어를 하나 더 추가했다. vgg가 아마 softmax로 1000개의 클래스를 구분하는 걸로 학습되어서 마지막 linear layer가 1000개이므로 이에 맞춘 것.

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


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

아니면 모델의 레이어가 이렇게 있을 때

(classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear
    )
vgg.classifier._modules['6'] = torch.nn.Linear(4096, 1)

모델의 레이어를 선택해서 바꿔줄 수도 있다.

freezing 하기

제일 마지막에 추가한 linear 레이어를 제외하고는 사실 학습된 부분이라 업데이트를 할 필요가 없다. 이를 위해서는 밑과 같이 grad를 꺼주면 된다.

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

반면 마지막 레이어는 필요하므로

for param in my_model.linear_layers.parameters():
    param.requires_grad = True
profile
AI Engineer
post-custom-banner

0개의 댓글