Visualizing Neural Network Structures in PyTorch

Click on the “CVer“, and choose to “star” or “top” it

Heavyweight content delivered at the first time

Author: Tian Haishan

https://zhuanlan.zhihu.com/p/66320870

This article is authorized, and no secondary reproduction is allowed without permission

Installation

You can install it using the following commands

conda install pytorch-nightly -c pytorch
conda install graphviz
conda install torchvision
conda install tensorwatch

This tutorial is based on the following versions:

torchvision.__version__   '0.2.1'
torch.__version__         '1.2.0.dev20190610'
sys.version               '3.6.8 |Anaconda custom (64-bit)| (default, Dec 30 2018, 01:22:34) \n[GCC 7.3.0]'

Load Libraries

import sys
import torch
import tensorwatch as tw
import torchvision.models

Visualizing Network Structure

alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224])

Load alexnet, the draw_model function requires three parameters: the first is the model, the second parameter is input_shape, and the third parameter is orientation, which can be ‘LR’ or ‘TB’, representing left-right layout and top-bottom layout respectively.

In the notebook, after executing the above code, the following figure will be displayed, visualizing the structure of the network and the names and shapes of each layer.

Statistics of Network Parameters

You can use the model_stats method to count the parameters of each layer.

tw.model_stats(alexnet_model, [1, 3, 224, 224])

[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!

alexnet_model.features

Sequential(
  (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
  (1): ReLU(inplace=True)
  (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
  (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (4): ReLU(inplace=True)
  (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
  (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (7): ReLU(inplace=True)
  (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (9): ReLU(inplace=True)
  (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (11): ReLU(inplace=True)
  (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
)

alexnet_model.classifier

Sequential(
  (0): Dropout(p=0.5)
  (1): Linear(in_features=9216, out_features=4096, bias=True)
  (2): ReLU(inplace=True)
  (3): Dropout(p=0.5)
  (4): Linear(in_features=4096, out_features=4096, bias=True)
  (5): ReLU(inplace=True)
  (6): Linear(in_features=4096, out_features=1000, bias=True)
)

Reference

https://github.com/microsoft/tensorwatch

CVer Academic Exchange Group

Scan to add CVer assistant, you can apply to join the CVer-Object Detection Group, Image Segmentation, Object Tracking, Face Detection & Recognition, OCR, Super Resolution, SLAM, Medical Imaging, Re-ID, and GAN groups. Be sure to note:Research direction + location + school/company + nickname (e.g., Object Detection + Shanghai + SJTU + Kaka)

▲ Long press to join the group

Such a hard content share, please give me a look

▲ Long press to follow us

Please give me a look!