Pytorch网络结构可视化：torchsummary

今天药忘吃喽~ 2023-09-27 04:55 94阅读 0赞

## 一、单通道输入网络 ##

单通道输入的情况大致有以下两种结构：

### 1、结构1：只有一条路可以走 ###

![在这里插入图片描述][e7298d1a1d8046b8b7dda533af4b3c51.png]

### 2、结构2：输入为一条路，输出为多条路 ###

![在这里插入图片描述][8df5fa12c8b249459b81211acf2c9a24.png]  
以上两种的输入只有一个input，这种是经常遇到的情况。

import torch
    import torch.nn as nn
    from torchsummary import summary
    
    class Network(nn.Module): 
        def __init__(self, channels_img, features_d, num_classes, img_size): 
            super(Network, self).__init__()
            self.img_size = img_size
            self.disc = nn.Conv2d(
                in_channels = channels_img + 1, 
                out_channels = features_d, 
                kernel_size = (4,4)
            )
    
            # ConditionalGan: 
            self.embed = nn.Embedding(
                num_embeddings = num_classes, 
                embedding_dim = img_size * img_size
            )
    
       def forward(self, x, labels): 
            embedding = self.embed(labels).view(labels.shape[0], 1, self.img_size, self.img_size)
            x = torch.cat([x, embedding], dim = 1)
            return self.disc(x) 
        
    # device: 
    device = torch.device("cpu")
    
    # hyperparameter: 
    batch_size = 64
    
    # Initialize model: 
    model = Network(
        channels_img = 1, 
        features_d = 16, 
        num_classes = 10, 
        img_size = 28).to(device) 
    
    # Print model summary: 
    summary(
        model, 
        input_size = [(1, 28, 28), (1, 28, 28)], # MNIST
        batch_size = batch_size
    )

--------------------

参考资料：  
[PYTORCH网络结构可视化（TORCHSUMMARY+TENSORBOARDX+网络多输通道入情况+多分支网络）][PYTORCH_TORCHSUMMARY_TENSORBOARDX]  
[在Python中嵌入层:如何在Torchsummary中正确使用?][Python_Torchsummary]  
[pytorch 网络结构可视化方法汇总（三种实现方法详解）][pytorch]

[e7298d1a1d8046b8b7dda533af4b3c51.png]: https://img-blog.csdnimg.cn/e7298d1a1d8046b8b7dda533af4b3c51.png
[8df5fa12c8b249459b81211acf2c9a24.png]: https://img-blog.csdnimg.cn/8df5fa12c8b249459b81211acf2c9a24.png
[PYTORCH_TORCHSUMMARY_TENSORBOARDX]: https://www.freesion.com/article/49401186740/
[Python_Torchsummary]: https://www.it1352.com/2294777.html
[pytorch]: https://blog.csdn.net/qq_27825451/article/details/96856217