feature map的可视化

最近使用yolo-tiny进行目标检测，需要用tensorboard将特征图可视化，从网上搜索了一大圈，找到了比较靠谱的方法。`

img_grid = vutils.make_grid(x, normalize=True, scale_each=True, nrow=2)
# 绘制原始图像
writer.add_image('raw img', img_grid, global_step=666)  # j 表示feature map数
print(x.size())

model.eval()
for name, layer in model._modules.items():

    # 为fc层预处理x
    x = x.view(x.size(0), -1) if "fc" in name else x
    print(x.size())

    x = layer(x)
    print(f'{name}')

    # 第一个卷积没有进行relu，要记得加上
    x = F.relu(x) if 'conv' in name else x
    if  'layer' in name or 'conv' in name:
        x1 = x.transpose(0, 1)  # C，B, H, W  ---> B，C, H, W
        img_grid = vutils.make_grid(x1, normalize=True, scale_each=True, nrow=4)  # normalize进行归一化处理
        writer.add_image(f'{name}_feature_maps', img_grid, global_step=0)

因为我需要把网络内部分节点的特征图可视化出来，特征图的每个像素点上的数值范围不是[0,1]，而是可正可负，可大可小，因此需要做一些特殊处理。这里就要用到 ==torchvision.utils.make_grid( )==函数，把输入的特征图做一个归一化，把参数normalize设置为True即可，它能帮我们把数据的输入范围调整至[0, 1]之间

def make_grid(tensor, nrow=8, padding=2,
              normalize=False, range=None, scale_each=False, pad_value=0):

上面这些方法是针对一些backbone比较简单直接的网络,比如vgg，对于yolo这种backbone比较复杂，一层套一层的，还是不能直接用的。拿yolo4-tiny举例子，backbone为：

#-------------------------------------------------#
#   卷积块
#   CONV+BATCHNORM+LeakyReLU
#-------------------------------------------------#
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super(BasicConv, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.activation = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activation(x)
        return x
#---------------------------------------------------#
#   CSPdarknet53-tiny的结构块
#   存在一个大残差边
#   这个大残差边绕过了很多的残差结构
#---------------------------------------------------#
class Resblock_body(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Resblock_body, self).__init__()
        self.out_channels = out_channels

        self.conv1 = BasicConv(in_channels, out_channels, 3)

        self.conv2 = BasicConv(out_channels//2, out_channels//2, 3)
        self.conv3 = BasicConv(out_channels//2, out_channels//2, 3)

        self.conv4 = BasicConv(out_channels, out_channels, 1)
        self.maxpool = nn.MaxPool2d([2,2],[2,2])

    def forward(self, x):
        x = self.conv1(x)
        route = x
        
        c = self.out_channels
        x = torch.split(x, c//2, dim=1)[1]
        x = self.conv2(x)
        route1 = x
        x = self.conv3(x)
        x = torch.cat([x,route1], dim = 1) 
        x = self.conv4(x)
        feat = x

        x = torch.cat([route, x], dim=1)
        x = self.maxpool(x)
        return x,feat

class CSPDarkNet(nn.Module):
    def __init__(self):
        super(CSPDarkNet, self).__init__()
        self.conv1 = BasicConv(3, 32, kernel_size=3, stride=2)
        self.conv2 = BasicConv(32, 64, kernel_size=3, stride=2)

        self.resblock_body1 =  Resblock_body(64, 64)
        self.resblock_body2 =  Resblock_body(128, 128)
        self.resblock_body3 =  Resblock_body(256, 256)
        self.conv3 = BasicConv(512, 512, kernel_size=3)


        self.num_features = 1
        # 进行权值初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x, _    = self.resblock_body1(x)
        x, _    = self.resblock_body2(x)
        x, feat1    = self.resblock_body3(x)
        x = self.conv3(x)
        feat2 = x
        return feat1,feat2

其中Resblock_body()的结构图是

看起来还是容易理解的。一般layer(image)的输入格式为tensor，经过简单的卷积层输出也格式为tensor()，但是经过Resblock_body()时输出的格式却为tuple(),这样就不能继续循环了。
最后将代码进行了粗糙的修改：

#-------------------------------------#
#       对单张图片进行预测
#-------------------------------------#
from yolo import YOLO
from PIL import Image
from nets.CSPdarknet53_tiny import Resblock_body
from torchvision.utils import make_grid
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F
from nets.yolo4_tiny import YoloBody
import  torch
import numpy as np

writer = SummaryWriter(log_dir='runs/3', flush_secs=60)


yolo = YOLO()

# while True:
    # image = input('Input image filename:')
image = "./img/street.jpg"

try:
    image = Image.open(image)
    image = image.convert('RGB')
except:
    print('Open Error! Try again!')
    # continue
else:

    # 绘制原始图像
    model = YoloBody(400, 400)
    model_path = "model_data/cosine.pth"
    # 加快模型训练的效率
    print('Loading weights into state dict...')
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model_dict = model.state_dict()
    pretrained_dict = torch.load(model_path, map_location=device)
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    print('Finished!')

    from torchvision import transforms
    loader = transforms.Compose([
        transforms.ToTensor()])
    # image = Image.open(image).convert('RGB')
    image = loader(image).unsqueeze(0)  # 将图片转化为tensor格式

    model.eval()
    for name, layer in model.backbone._modules.items():
        # 为fc层预处理x
        # images = images.view(images.size(0), -1) if "fc" in name else images
        # print(image.size())
        if name == "resblock_body1":
            image = layer.conv1(image)
            route = image
            image = torch.split(image, 32, dim=1)[1]
            image = layer.conv2(image)
            route1 = image
            image = layer.conv3(image)
            image = torch.cat([image, route1], dim=1)
            image = layer.conv4(image)
            image = torch.cat([route, image], dim=1)
            image = layer.maxpool(image)[0].unsqueeze(0)
        else:
            if name == "resblock_body2":
                image = layer.conv1(image)
                route = image
                image = torch.split(image, 64, dim=1)[1]
                image = layer.conv2(image)
                route1 = image
                image = layer.conv3(image)
                image = torch.cat([image, route1], dim=1)
                image = layer.conv4(image)
                image = torch.cat([route, image], dim=1)
                image = layer.maxpool(image)[0].unsqueeze(0)
            else:
                if name == "resblock_body3":
                    image = layer.conv1(image)
                    route = image
                    image = torch.split(image, 128, dim=1)[1]
                    image = layer.conv2(image)
                    route1 = image
                    image = layer.conv3(image)
                    image = torch.cat([image, route1], dim=1)
                    image = layer.conv4(image)
                    image = torch.cat([route, image], dim=1)
                    image = layer.maxpool(image)[0].unsqueeze(0)
                else:
                    image = layer(image)
        print(f'{name}')

        # 第一个卷积没有进行relu，要记得加上
        if 'layer' in name or 'conv' in name:
            x1 = image.transpose(0, 1)  # C，B, H, W  ---> B，C, H, W
            img_grid = make_grid(x1, normalize=True, scale_each=True, nrow=4)  # normalize进行归一化处理
            writer.add_image(f'{name}_feature_maps', img_grid, global_step=0, dataformats='CHW')

if循环中又臭又长的代码就是Resblock_body()的内容，只能这样一行一行的抠出来。
最后展示效果