GoogLeNet で出力層の手前の層を特徴ベクトルとして取得

GoogLeNet で出力層の手前の層を特徴ベクトルとして取得する方法のメモ。
register_forward_hook() で出力層の手前の層の出力を取得します。

# -*- coding:utf-8 -*-

import sys
import json
import torch
import torchvision.transforms as transforms
from PIL import Image
from googlenet_pytorch import GoogLeNet

feature_vector = None

def get_feature_vector(preproc, model, img_file):
    input_image = Image.open(img_file)
    input_tensor = preproc(input_image)
    input_batch = input_tensor.unsqueeze(0)
    
    logits = model(input_batch)
    preds = torch.topk(logits, k=5).indices.squeeze(0).tolist()
    return feature_vector

def forward_hook(module, inputs, outputs):
    global feature_vector
    feature_vector = outputs.detach().clone()[0].tolist()
    
def init():
    preproc = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])

    labels_map = json.load(open("labels_map.txt"))
    labels_map = [labels_map[str(i)] for i in range(1000)]

    model = GoogLeNet.from_pretrained("googlenet")
    model.eval()

    layers = list(model.children())
    handle = layers[-2].register_forward_hook(forward_hook)
    return preproc, model

def main():
    preproc, model = init()
    img_file = sys.argv[1]
    fv = get_feature_vector(preproc, model, img_file)
    print(fv)

    return 0

if __name__ == '__main__':
    res = main()
    exit(res)