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))
저장하기
torch.save(model.state_dict(), "model.pt"))
불러오기
new_model.load_state_dict(torch.load("model.pt")))
저장하기
torch.save(model,"model_pickle.pt"))
불러오기
model = torch.load("model_pickle.pt"))
pretrained model 활용 시 모델의 일부분은 frozen 시킴
torchvision에서 모델을 가져올 수 있는데 이번에는 vgg 모델을 가져와보겠다. pretrained 옵션을 줘야 w도 같이 온다.
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)
모델의 레이어를 선택해서 바꿔줄 수도 있다.
제일 마지막에 추가한 linear 레이어를 제외하고는 사실 학습된 부분이라 업데이트를 할 필요가 없다. 이를 위해서는 밑과 같이 grad를 꺼주면 된다.
for param in my_model.parameters():
param.requires_grad = False
반면 마지막 레이어는 필요하므로
for param in my_model.linear_layers.parameters():
param.requires_grad = True