利用Pytorch實現獲取特徵圖的方法詳解

2022-10-16 14:01:51

簡單載入官方預訓練模型

torchvision.models預定義了很多公開的模型結構

如果pretrained引數設定為False，那麼僅僅設定模型結構；如果設定為True，那麼會啟動一個下載流程，下載預訓練引數

如果只想呼叫模型，不想訓練，那麼設定model.eval()和model.requires_grad_(False)

想檢視模型引數可以使用modules和named_modules，其中named_modules是一個長度為2的tuple，第一個變數是name，第二個變數是module本身。

# -*- coding: utf-8 -*-
from torch import nn
from torchvision import models

# load model. If pretrained is True, there will be a downloading process
model = models.vgg19(pretrained=True)
model.eval()
model.requires_grad_(False)

# get model component
features = model.features
modules = features.modules()
named_modules = features.named_modules()

# print modules
for module in modules:
    if isinstance(module, nn.Conv2d):
        weight = module.weight
        bias = module.bias
        print(module, weight.shape, bias.shape,
              weight.requires_grad, bias.requires_grad)
    elif isinstance(module, nn.ReLU):
        print(module)

print()
for named_module in named_modules:
    name = named_module[0]
    module = named_module[1]
    if isinstance(module, nn.Conv2d):
        weight = module.weight
        bias = module.bias
        print(name, module, weight.shape, bias.shape,
              weight.requires_grad, bias.requires_grad)
    elif isinstance(module, nn.ReLU):
        print(name, module)

圖片預處理

使用opencv和pil讀圖都可以使用transforms.ToTensor()把原本[H, W, 3]的資料轉成[3, H, W]的tensor。但opencv要注意把資料改成RGB順序。

vgg系列模型需要做normalization，建議配合torchvision.transforms來實現。

mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W are expected to be at least 224. The images have to be loaded in to a range of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225].

參考：https://pytorch.org/hub/pytorch_vision_vgg/

# -*- coding: utf-8 -*-
from PIL import Image
import cv2
import torch
from torchvision import transforms

# transforms for preprocess
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# load image using cv2
image_cv2 = cv2.imread('lena_std.bmp')
image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
image_cv2 = preprocess(image_cv2)

# load image using pil
image_pil = Image.open('lena_std.bmp')
image_pil = preprocess(image_pil)

# check whether image_cv2 and image_pil are same
print(torch.all(image_cv2 == image_pil))
print(image_cv2.shape, image_pil.shape)

提取單個特徵圖

如果只提取單層特徵圖，可以把模型截斷，以節省算力和視訊記憶體消耗。

下面索引之所以有+1是因為pytorch預訓練模型裡面第一個索引的module總是完整模組結構，第二個才開始子模組。

# -*- coding: utf-8 -*-
from PIL import Image
from torchvision import models
from torchvision import transforms

# load model. If pretrained is True, there will be a downloading process
model = models.vgg19(pretrained=True)
model = model.features[:16 + 1]  # 16 = conv3_4
model.eval()
model.requires_grad_(False)
model.to('cuda')
print(model)

# load and preprocess image
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
    transforms.Resize(size=(224, 224))
])
image = Image.open('lena_std.bmp')
image = preprocess(image)
inputs = image.unsqueeze(0)  # add batch dimension
inputs = inputs.cuda()

# forward
output = model(inputs)
print(output.shape)

提取多個特徵圖

第一種方式：逐層執行model，如果碰到了需要儲存的feature map就存下來。

第二種方式：使用register_forward_hook，使用這種方式需要用一個類把feature map以成員變數的形式快取下來。

兩種方式的執行效率差不多

第一種方式簡單直觀，但是隻能處理類似VGG這種沒有跨層連線的網路；第二種方式更加通用。

# -*- coding: utf-8 -*-
from PIL import Image
import torch
from torchvision import models
from torchvision import transforms

# load model. If pretrained is True, there will be a downloading process
model = models.vgg19(pretrained=True)
model = model.features[:16 + 1]  # 16 = conv3_4
model.eval()
model.requires_grad_(False)
model.to('cuda')

# check module name
for named_module in model.named_modules():
    name = named_module[0]
    module = named_module[1]
    print('-------- %s --------' % name)
    print(module)
    print()

# load and preprocess image
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
    transforms.Resize(size=(224, 224))
])
image = Image.open('lena_std.bmp')
image = preprocess(image)
inputs = image.unsqueeze(0)  # add batch dimension
inputs = inputs.cuda()

# forward - 1
layers = [2, 7, 8, 9, 16]
layers = sorted(set(layers))
feature_maps = {}
feature = inputs
for i in range(max(layers) + 1):
    feature = model[i](feature)
    if i in layers:
        feature_maps[i] = feature
for key in feature_maps:
    print(key, feature_maps.get(key).shape)


# forward - 2
class FeatureHook:
    def __init__(self, module):
        self.inputs = None
        self.output = None
        self.hook = module.register_forward_hook(self.get_features)

    def get_features(self, module, inputs, output):
        self.inputs = inputs
        self.output = output


layer_names = ['2', '7', '8', '9', '16']
hook_modules = []
for named_module in model.named_modules():
    name = named_module[0]
    module = named_module[1]
    if name in layer_names:
        hook_modules.append(module)

hooks = [FeatureHook(module) for module in hook_modules]
output = model(inputs)
features = [hook.output for hook in hooks]
for feature in features:
    print(feature.shape)

# check correctness
for i, layer in enumerate(layers):
    feature1 = feature_maps.get(layer)
    feature2 = features[i]
    print(torch.all(feature1 == feature2))

使用第二種方式（register_forward_hook），resnet特徵圖也可以順利拿到。

而由於resnet的model已經不可以用model[i]的形式索引，所以無法使用第一種方式。

# -*- coding: utf-8 -*-
from PIL import Image
from torchvision import models
from torchvision import transforms

# load model. If pretrained is True, there will be a downloading process
model = models.resnet18(pretrained=True)
model.eval()
model.requires_grad_(False)
model.to('cuda')

# check module name
for named_module in model.named_modules():
    name = named_module[0]
    module = named_module[1]
    print('-------- %s --------' % name)
    print(module)
    print()

# load and preprocess image
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
    transforms.Resize(size=(224, 224))
])
image = Image.open('lena_std.bmp')
image = preprocess(image)
inputs = image.unsqueeze(0)  # add batch dimension
inputs = inputs.cuda()


class FeatureHook:
    def __init__(self, module):
        self.inputs = None
        self.output = None
        self.hook = module.register_forward_hook(self.get_features)

    def get_features(self, module, inputs, output):
        self.inputs = inputs
        self.output = output


layer_names = [
    'conv1',
    'layer1.0.relu',
    'layer2.0.conv1'
]

hook_modules = []
for named_module in model.named_modules():
    name = named_module[0]
    module = named_module[1]
    if name in layer_names:
        hook_modules.append(module)

hooks = [FeatureHook(module) for module in hook_modules]
output = model(inputs)
features = [hook.output for hook in hooks]
for feature in features:
    print(feature.shape)

問題來了，resnet這種型別的網路結構怎麼截斷？

使用如下命令就可以，print檢視需要截斷到哪裡，然後用nn.Sequential重組即可。

需注意重組後網路的module_name會發生變化。

print(list(model.children())
model = torch.nn.Sequential(*list(model.children())[:6])

以上就是利用Pytorch實現獲取特徵圖的方法詳解的詳細內容，更多關於Pytorch獲取特徵圖的資料請關注it145.com其它相關文章！

利用Pytorch實現獲取特徵圖的方法詳解

目錄

簡單載入官方預訓練模型

圖片預處理

提取單個特徵圖

提取多個特徵圖

熱門文章