pytorch对神经网络的特征图的可视-结合matplotlib工具

在前向传播中返回需要可视化的特征图：

def forward(self, x):
    output=[]
    layers = self.dla(x)

    p3 = self.p3(layers[3])
    p3 = self.p3_l2(p3)
    output.append(p3)

    p4 = self.p4(layers[4])
    p4 = self.p4_l2(p4)
    output.append(p4)

    p5 = self.p5(layers[5])
    p5 = self.p5_l2(p5)
    output.append(p5)
    return output

对返回的特征图进行显示

for feature in features_map:
    im=np.squeeze(feature.cpu().detach().numpy())
    im=np.transpose(im,[1,2,0])
    plt.figure()
    for i in range(16):
        ax=plt.subplot(4,4,i+1)
        plt.imshow(im[:,:,i])
    plt.show()

结果展示

原图片：
特征图：

tensorboard显示网络结构及train_loss

from tensorboardX import SummaryWriter
writer=SummaryWriter(log_dir='logs',comment='l1loss')
......
graph_inputs=torch.from_numpy(np.random.rand(1,1,28,28)).type(torch.FloatTensor).cuda()
writer.add_graph(model,graph_inputs) //加网络结构图
writer.add_scalar('Train_loss',train_loss,epoch) //加trainloss 第三个参数是横坐标

然后在cmd中输入tensorboard --logdir=路径