PyTorch-模型可视化工具TorchSummary

Posted 周先森爱吃素

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了PyTorch-模型可视化工具TorchSummary相关的知识,希望对你有一定的参考价值。

简介

不同于TensorboardX对Tensorboard的支持以方便了PyTorch的训练可视化,PyTorch并没有很好的模型可视化工具,TorchSummary对此做出了补足,极大降低了模型可视化难度,也方便模型参数等数据的统计。本文介绍TorchSummary这个小工具的使用。

安装

使用pip安装即可。

pip install torchsummary

开发缘由

首先,我们知道,PyTorch其实自带模型可视化的功能,其基础调用格式如下。

print(model)
import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

例如,下述简单模型通过print可视化,结果如下,显然,这只是对模型含有的modules做了一个对象及其参数打印,我们更希望输出每一层的layer类型、参数量以及输出feature map尺寸等。

Net(
  (conv1): Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1))
  (conv2_drop): Dropout2d(p=0.5, inplace=False)
  (fc1): Linear(in_features=320, out_features=50, bias=True)
  (fc2): Linear(in_features=50, out_features=10, bias=True)
)

对此,TorchSummary提供了更详细的信息分析,包括模块信息(每一层的类型、输出shape和参数量)、模型整体的参数量、模型大小、一次前向或者反向传播需要的内存大小等

使用教程

TorchSummary的使用基于下述核心API,只要提供给summary函数模型以及输入的size就可以了。

from torchsummary import summary
summary(model, input_size=(channels, H, W))

如在一个简单CNN上进行模型可视化,代码和结果如下(测试均使用PyTorch1.6.0),可视化输出包括我上一节文末提到的我们需要的常用信息,非常丰富。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net().to(device)

summary(model, (1, 28, 28))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------

对于多输入的情况,只要传入的input_size改为一个安装输入所需size组成的列表就行,示例如下。

import torch
import torch.nn as nn
from torchsummary import summary


class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )

    def forward(self, x, y):
        x1 = self.features(x)
        x2 = self.features(y)
        return x1, x2


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleConv().to(device)

summary(model, input_size=[(1, 16, 16), (1, 28, 28)])

结果如下。

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 1, 16, 16]              10
              ReLU-2            [-1, 1, 16, 16]               0
            Conv2d-3            [-1, 1, 28, 28]              10
              ReLU-4            [-1, 1, 28, 28]               0
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.77
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.78
----------------------------------------------------------------

补充说明

本文简单介绍了我比较喜欢的PyTorch模型可视化工具,文中示例代码参考官网,如果对你有所帮助,麻烦点赞支持一下。

以上是关于PyTorch-模型可视化工具TorchSummary的主要内容,如果未能解决你的问题,请参考以下文章

PyTorch 打印模型结构输出维度和参数信息(torchsummary)

PyTorch打印模型结构输出维度和参数信息(torchsummary)

Tensorflow,pytorch查看模型参数,模型可视化

PyTorch教程-8:举例详解TensorBoard的使用

Model Log模型评估指标可视化,自动画LossAccuracy曲线图工具,无需人工参与!

[十九]深度学习Pytorch-可视化工具TensorBoard