使用pytorch提取折積神經網路的特徵圖視覺化

2022-03-29 13:03:32

前言

文章中的程式碼是參考基於Pytorch的特徵圖提取編寫的程式碼本身很簡單這裡只做簡單的描述。

1. 效果圖

先看效果圖(第一張是原圖，後面的都是相應的特徵圖，這裡使用的網路是resnet50，需要注意的是下面圖片顯示的特徵圖是經過放大後的圖，原圖是比較小的圖，因為太小不利於我們觀察):

2. 完整程式碼

import os
import torch
import torchvision as tv
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import argparse
import skimage.data
import skimage.io
import skimage.transform
import numpy as np
import matplotlib.pyplot as plt
import torchvision.models as models
from PIL import Image
import cv2

class FeatureExtractor(nn.Module):
    def __init__(self, submodule, extracted_layers):
        super(FeatureExtractor, self).__init__()
        self.submodule = submodule
        self.extracted_layers = extracted_layers
 
    def forward(self, x):
        outputs = {}
        for name, module in self.submodule._modules.items():
            if "fc" in name: 
                x = x.view(x.size(0), -1)
            
            x = module(x)
            print(name)
            if self.extracted_layers is None or name in self.extracted_layers and 'fc' not in name:
                outputs[name] = x

        return outputs


def get_picture(pic_name, transform):
    img = skimage.io.imread(pic_name)
    img = skimage.transform.resize(img, (256, 256))
    img = np.asarray(img, dtype=np.float32)
    return transform(img)

def make_dirs(path):
    if os.path.exists(path) is False:
        os.makedirs(path)


def get_feature():
    pic_dir = './images/2.jpg'
    transform = transforms.ToTensor()
    img = get_picture(pic_dir, transform)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 插入維度
    img = img.unsqueeze(0)

    img = img.to(device)

    
    net = models.resnet101().to(device)
    net.load_state_dict(torch.load('./model/resnet101-5d3b4d8f.pt'))
    exact_list = None
    dst = './feautures'
    therd_size = 256

    myexactor = FeatureExtractor(net, exact_list)
    outs = myexactor(img)
    for k, v in outs.items():
        features = v[0]
        iter_range = features.shape[0]
        for i in range(iter_range):
            #plt.imshow(x[0].data.numpy()[0,i,:,:],cmap='jet')
            if 'fc' in k:
                continue

            feature = features.data.numpy()
            feature_img = feature[i,:,:]
            feature_img = np.asarray(feature_img * 255, dtype=np.uint8)
            
            dst_path = os.path.join(dst, k)
            
            make_dirs(dst_path)
            feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)
            if feature_img.shape[0] < therd_size:
                tmp_file = os.path.join(dst_path, str(i) + '_' + str(therd_size) + '.png')
                tmp_img = feature_img.copy()
                tmp_img = cv2.resize(tmp_img, (therd_size,therd_size), interpolation =  cv2.INTER_NEAREST)
                cv2.imwrite(tmp_file, tmp_img)
            
            dst_file = os.path.join(dst_path, str(i) + '.png')
            cv2.imwrite(dst_file, feature_img)

if __name__ == '__main__':
    get_feature()

3. 程式碼說明

下面的模組是根據所指定的模型篩選出指定層的特徵圖輸出，如果未指定也就是extracted_layers是None則以字典的形式輸出全部的特徵圖，另外因為全連線層本身是一維的沒必要輸出因此進行了過濾。

class FeatureExtractor(nn.Module):
    def __init__(self, submodule, extracted_layers):
        super(FeatureExtractor, self).__init__()
        self.submodule = submodule
        self.extracted_layers = extracted_layers
 
    def forward(self, x):
        outputs = {}
        for name, module in self.submodule._modules.items():
            if "fc" in name: 
                x = x.view(x.size(0), -1)
            
            x = module(x)
            print(name)
            if self.extracted_layers is None or name in self.extracted_layers and 'fc' not in name:
                outputs[name] = x

        return outputs

這段主要是儲存圖片，為每個層建立一個資料夾將特徵圖以JET的colormap進行按順序儲存到該資料夾，並且如果特徵圖過小也會對特徵圖放大同時儲存原始圖和放大後的圖。

for k, v in outs.items():
        features = v[0]
        iter_range = features.shape[0]
        for i in range(iter_range):
            #plt.imshow(x[0].data.numpy()[0,i,:,:],cmap='jet')
            if 'fc' in k:
                continue

            feature = features.data.numpy()
            feature_img = feature[i,:,:]
            feature_img = np.asarray(feature_img * 255, dtype=np.uint8)
            
            dst_path = os.path.join(dst, k)
            
            make_dirs(dst_path)
            feature_img = cv2.applyColorMap(feature_img, cv2.COLORMAP_JET)
            if feature_img.shape[0] < therd_size:
                tmp_file = os.path.join(dst_path, str(i) + '_' + str(therd_size) + '.png')
                tmp_img = feature_img.copy()
                tmp_img = cv2.resize(tmp_img, (therd_size,therd_size), interpolation =  cv2.INTER_NEAREST)
                cv2.imwrite(tmp_file, tmp_img)
            
            dst_file = os.path.join(dst_path, str(i) + '.png')
            cv2.imwrite(dst_file, feature_img)

這裡主要是一些引數，比如要提取的網路，網路的權重，要提取的層，指定的影象放大的大小，儲存路徑等等。

	net = models.resnet101().to(device)
    net.load_state_dict(torch.load('./model/resnet101-5d3b4d8f.pt'))
    exact_list = None#['conv1']
    dst = './feautures'
    therd_size = 256

4. 視覺化梯度，feature

上面的辦法只是簡單的將經過網路計算的圖片的輸出的feature進行圖片，github上有將CNN的梯度等全部進行視覺化的程式碼：pytorch-cnn-visualizations，需要注意的是如果只是簡單的替換成自己的網路可能無法執行，大概率會報model沒有features或者classifier等錯誤，這兩個是進行分類網路定義時的Sequential，其實就是索引網路的每一層，自己稍微修改用model.children()等方法進行替換即可，我自己修改之後得到的程式碼grayondream-pytorch-visualization（本來想稍微封裝一下成為一個更加通用的結構，暫時沒時間以後再說吧！），下面是效果圖：