torchsummary用法

weixin_50508552

已于 2024-06-05 20:43:46 修改

阅读量2.3k

点赞数 26

文章标签：深度学习人工智能

于 2024-06-05 19:59:44 首次发布

本文链接：https://blog.csdn.net/weixin_50508552/article/details/139480906

版权

torchsummary 是一个用于 PyTorch 模型的概述工具，可以方便地查看模型的结构、参数数量以及每一层的输出形状。以下是如何使用 torchsummary 的详细步骤：

import torch
import torch.nn as nn
import torchsummary

class SimpleModel(nn.Module):
def __init__(self):
	super(SimpleModel, self).__init__()
	self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
	self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
	self.fc1 = nn.Linear(32 * 8 * 8, 128)
	self.fc2 = nn.Linear(128, 10)
	

def forward(self, x):
    x = self.conv1(x)
    x = torch.relu(x)
    x = torch.max_pool2d(x, 2)
    x = self.conv2(x)
    x = torch.relu(x)
    x = torch.max_pool2d(x, 2)
    x = x.view(x.size(0), -1)  # flatten the tensor
    x = self.fc1(x)
    x = torch.relu(x)
    x = self.fc2(x)
    return x

model = SimpleModel()
torchsummary.summary(model, input_size=(3, 32, 32))

可以输出模型每一层的输出形状、参数信息。

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 16, 32, 32]             448
            Conv2d-2           [-1, 32, 16, 16]           4,640
            Linear-3                  [-1, 128]         262,272
            Linear-4                   [-1, 10]           1,290
================================================================
Total params: 268,650
Trainable params: 268,650
Non-trainable params: 0
----------------------------------------------------------------

在 torchsummary 的输出中，Output Shape 列显示了每一层输出的张量的形状。

其中：

Output Shape 列中每一项的格式是 [-1, C, H, W] 或者 [-1, D]。每个维度的具体含义如下：

-1: 这个位置代表批量大小（batch size），在实际使用中是动态的，因此通常用 -1 表示。这个值在模型定义中不会固定，因为每次输入的批量大小可以不同。
C: 通道数（channels）。对于卷积层输出，C 表示输出特征图的通道数；对于全连接层输出，C 表示输出的特征维度。
H: 特征图的高度（height）。对于卷积层和池化层，H 是输出特征图的高度。对于全连接层（Linear），这个维度不存在。
W: 特征图的宽度（width）。对于卷积层和池化层，W 是输出特征图的宽度。对于全连接层（Linear），这个维度不存在。

对于全连接层（Linear），输出形状通常是 [-1, D]，其中 D 表示输出的特征维度（单个向量的长度）。

这一个例子弄明白！！！！

1. 首先，第一个参数-1不用管。指batch_size

2. 整个model的输入大小为input_size=(3, 32, 32)

3. conv1层，in_channel数为3，因为输入的参数的通道数就是3

输出的通道数out_channel数为16。

输出的图像大小，即output_shape后两个参数，计算方式为：

Hout = （Hin+2*padding-kernel)/stride + 1

W同理

所以输出大小h= （32+2*1-3）/1 + 1 = 32

所以conv1的输出为[-1,16,32,32]

4. max_pool1层，承接在conv1层后面，且