网络结构可视化——Torch-summary库

Torchsummary库是深度学习网络结构可视化常用的库：安装地址
Torch-summary库是torchsummary的加强版，库的介绍和安装地址.
建议安装Torch-summary库而非Torchsummary库，前者在继承后者的函数外还解决了后者存在的诸多Bug

Torchsummary库常遇问题

一、问题一：使用torchsummary查看网络结构时报错：AttributeError: ‘list’ object has no attribute ‘size’

解决方法：

pip uninstall torchsummary        # 卸载原来的torchsummary库
pip install torch-summary==1.4.4  # 安装升级版本torch-summary

例子：

from torchvision import models
import torchsummary as summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)
summary(vgg,(3,224,224))

问题二：torchsummary报错：TypeError: ‘module’ object is not callable

解决方案：

from torchvision import models
import torchsummary as summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)
summary.summary(vgg,(3,224,224))

Torch-summary库常见用法

# 使用样式
from torchsummary import summary
summary(model, input_size=(channels, H, W))

# 多输入情况并且打印不同层的特征图大小
from torchsummary import summary
summary(model,first_input,second_input)

# 打印不同的内容
import torch
import torch.nn as nn
from torch-summary import summary

class LSTMNet(nn.Module):
    """ Batch-first LSTM model. """
    def __init__(self, vocab_size=20, embed_dim=300, hidden_dim=512, num_layers=2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.decoder = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        embed = self.embedding(x)
        out, hidden = self.encoder(embed)
        out = self.decoder(out)
        out = out.view(-1, out.size(2))
        return out, hidden

summary(
    LSTMNet(),
    (100,),
    dtypes=[torch.long],
    branching=False,
    verbose=2,
    col_width=16,
    col_names=["kernel_size", "output_size", "num_params", "mult_adds"],)

打印结果：