[PyTorch] 모델 중간 레이어에서 특징 추출하기(get the intermediate features from the model)

 이번 포스팅에서는 모델의 중간 레이어에서 특징을 추출하는 방법을 살펴보겠습니다. ㅎㅎ

 

 pretrained VGG19을 불러오고, 파라미터를 freeze 합니다.

# pretrained VGG19를 불러옵니다
import torchvision.models as models

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_vgg = models.vgg19(pretrained=True).features.to(device).eval()

# 파라미터를 freeze 합니다.
for param in model_vgg.parameters():
    param.requires_grad_(False)

 

 모델의 중간 레이어의 특징을 얻는 함수를 정의합니다.

# 모델의 중간 레이어의 출력값을 얻는 함수를 정의합니다.
def get_features(x, model, layers):
    features = {}
    for name, layer in enumerate(model.children()): # 0, conv
        x = layer(x)
        if str(name) in layers:
            features[layers[str(name)]] = x
    return features

 

 model.children을 enumerate로 호출하면, 연산 순서와 연산을 얻을 수 있습니다. 차례대로 연산을 수행하며, layers로 원하는 연산에 해당하는 순서를 지정하여 입력하면 해당 순서의 연산의 출력값을 얻을 수 있습니다.

 

feature_layers = {'0': 'conv1_1',
                  '5': 'conv2_1',
                  '10': 'conv3_1',
                  '19': 'conv4_1',
                  '21': 'conv4_2',
                  '28': 'conv5_1'}
                  
con_tensor = content_tensor.unsqueeze(0).to(device) # 3,128,128
content_features = get_features(con_tensor, model_vgg, feature_layers)

 

 VGG19 Net에서 0,5,10,19,21,28 번째 연산은 conv연산입니다. 위 명령어로 해당하는 연산의 출력값을 얻을 수 있습니다. 감사합니다.