显示 torch.tensor 图像卷积特征图

参考博文：https://blog.csdn.net/dcrmg/article/details/81255498/

1.定义特征显示函数

def get_row_col(num_pic):
    squr = num_pic ** 0.5
    row = round(squr)
    col = row + 1 if squr - row > 0 else row
    return row, col


def visualize_feature_map(img_batch):
    img_batch = img_batch.detach().numpy()
    feature_map = np.squeeze(img_batch, axis=0)
    print(feature_map.shape)

    feature_map_combination = []
    plt.figure()

    num_pic = feature_map.shape[0]
    row, col = get_row_col(num_pic)

    for i in range(0, num_pic):
        feature_map_split = feature_map[i, :, :]
        feature_map_combination.append(feature_map_split)
        plt.subplot(row, col, i + 1)
        plt.imshow(feature_map_split)
        axis('off')
        # title('feature_map_{}'.format(i))

    plt.savefig('feature_map.png')
    plt.show()

    # 各个特征图按1：1 叠加
    feature_map_sum = sum(ele for ele in feature_map_combination)
    plt.imshow(feature_map_sum)
    plt.savefig("feature_map_sum.png")

2.定义一个卷积提取网络（博文里面使用的是keras,我这边改成了pytorch）

def create_model(use_bias=True):
    model = nn.Sequential()

    # 第一层CNN
    model.append(nn.Conv2d(3, 64,kernel_size=3,stride=1, padding=1,bias=use_bias))
    model.append(nn.BatchNorm2d(64))
    model.append(nn.LeakyReLU(0.2, True))
    model.append(nn.MaxPool2d(kernel_size=2, stride=2))

    # 第二层CNN
    model.append(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=use_bias))
    model.append(nn.BatchNorm2d(64))
    model.append(nn.LeakyReLU(0.2, True))
    model.append(nn.MaxPool2d(kernel_size=2, stride=2))

    # 第三层CNN
    model.append(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=use_bias))
    model.append(nn.BatchNorm2d(64))
    model.append(nn.LeakyReLU(0.2, True))

    # 第三层CNN
    model.append(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=use_bias))
    model.append(nn.BatchNorm2d(64))
    model.append(nn.LeakyReLU(0.2, True))

    return model

def forward(self, inputs):
    x = model(inputs)
    return x

3. 运行实例

先贴一张图

代码：图像格式要转一下，否则会出错

# coding: utf-8
import torch
from PIL import Image
from pylab import *
import torch.nn as nn

if __name__ == "__main__":
    #img = cv2.imread('001.jpg')
    img = Image.open('001.jpg')
    img = np.array(img)
    img = img.transpose(2, 0, 1)
    img_tensor = torch.from_numpy(img).float()
    img_tensor = torch.unsqueeze(img_tensor, 0)
    
    model = create_model()
    model.eval()
    
    # img_batch = np.expand_dims(img, axis=0)
    conv_img = model(img_tensor)  # conv_img 卷积结果
    visualize_feature_map(conv_img)

运行结果：