pretrained 모델에 fc레이어 추가해보기

한번도 모델에 추가해본적 없는데 한번 직접 해보려고 한다.

import torchvision.models as models
import torch.nn as nn
import torch
model = models.resnet18(pretrained = True)

num_features = model.fc.out_features
new_fc = nn.Linear(num_features, 10)
model.fc = nn.Sequential(
    model.fc, new_fc
)
print(model.fc)

Sequential(
  (0): Linear(in_features=512, out_features=1000, bias=True)
  (1): Linear(in_features=1000, out_features=10, bias=True)
)

이렇게 마지막 단에 10개의 클래스 변경한 것을 알 수 있다.

 

이번에는 여러개의 레이어도 추가해본다.

num_features = model.fc.out_features
fc_layers = [
    nn.Linear(num_features, 512),
    nn.ReLU(),
    nn.Linear(512, 10)
]

model.fc = nn.Sequential(model.fc,
                         *fc_layers)
print(model.fc)

Sequential(
  (0): Linear(in_features=512, out_features=1000, bias=True)
  (1): Linear(in_features=1000, out_features=512, bias=True)
  (2): ReLU()
  (3): Linear(in_features=512, out_features=10, bias=True)

이렇게 따로 추가한 것을 알 수 있다.

 

만약 layer를 추가하고 추가된 레이어만 gradient가 흐르게 하고 싶으면 다음과 같다.

model = models.resnet18(pretrained = True)

num_features = model.fc.out_features
fc_layers = [
    nn.Linear(num_features, 512),
    nn.ReLU(),
    nn.Linear(512, 10)
]

model.fc = nn.Sequential(model.fc,
                         *fc_layers)

for param in model.fc[:-1].parameters():
    param.requires_grad = False