快捷導(dǎo)航

使用pytorch提取卷積神經(jīng)網(wǎng)絡(luò)的特征圖可視化

更新時間：2022年03月29日 11:20:22 作者：落櫻彌城

這篇文章主要給大家介紹了關(guān)于使用pytorch提取卷積神經(jīng)網(wǎng)絡(luò)的特征圖可視化的相關(guān)資料,文中給出了詳細(xì)的思路以及示例代碼,對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友可以參考下

前言

文章中的代碼是參考基于Pytorch的特征圖提取編寫的代碼本身很簡單這里只做簡單的描述。

1. 效果圖

先看效果圖(第一張是原圖，后面的都是相應(yīng)的特征圖，這里使用的網(wǎng)絡(luò)是resnet50，需要注意的是下面圖片顯示的特征圖是經(jīng)過放大后的圖，原圖是比較小的圖，因為太小不利于我們觀察):

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. 代碼說明

下面的模塊是根據(jù)所指定的模型篩選出指定層的特征圖輸出，如果未指定也就是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

這段主要是存儲圖片，為每個層創(chuàng)建一個文件夾將特征圖以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)

這里主要是一些參數(shù)，比如要提取的網(wǎng)絡(luò)，網(wǎng)絡(luò)的權(quán)重，要提取的層，指定的圖像放大的大小，存儲路徑等等。

	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

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