print输出 pytorch_PyTorch如何打印每层的输出形状

最新推荐文章于 2024-10-23 17:28:13 发布

唐刘

最新推荐文章于 2024-10-23 17:28:13 发布

阅读量1.8k

点赞数 1

文章标签： print输出 pytorch

本文链接：https://blog.csdn.net/weixin_36143793/article/details/111890950

版权

本文展示了如何在PyTorch中使用torchsummary库来打印CNN（如MNSIT网络）和预训练模型（如VGG16）的每一层输出形状，包括输入输出尺寸和参数数量。

摘要由CSDN通过智能技术生成

安装：

pip install torchsummary or

git clone https://github.com/sksq96/pytorch-summary

使用示例：

CNN for MNSIT

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") # PyTorch v0.4.0

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

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

VGG16

import torch

from torchvision import models

from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

vgg = models.vgg16().to(device)

summary(vgg, (3, 224, 224))

>>>>>:

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

Layer (type) Output Shape Param #

================================================================

Conv2d-1 [-1, 64, 224, 224] 1,792

ReLU-2 [-1, 64, 224, 224] 0

Conv2d-3 [-1, 64, 224, 224] 36,928

ReLU-4 [-1, 64, 224, 224] 0

MaxPool2d-5 [-1, 64, 112, 112] 0

Conv2d-6 [-1, 128, 112, 112] 73,856

ReLU-7 [-1, 128, 112, 112] 0

Conv2d-8 [-1, 128, 112, 112] 147,584

ReLU-9 [-1, 128, 112, 112] 0

MaxPool2d-10 [-1, 128, 56, 56] 0

Conv2d-11 [-1, 256, 56, 56] 295,168

ReLU-12 [-1, 256, 56, 56] 0

Conv2d-13 [-1, 256, 56, 56] 590,080

ReLU-14 [-1, 256, 56, 56] 0

Conv2d-15 [-1, 256, 56, 56] 590,080

ReLU-16 [-1, 256, 56, 56] 0

MaxPool2d-17 [-1, 256, 28, 28] 0

Conv2d-18 [-1, 512, 28, 28] 1,180,160

ReLU-19 [-1, 512, 28, 28] 0

Conv2d-20 [-1, 512, 28, 28] 2,359,808

ReLU-21 [-1, 512, 28, 28] 0

Conv2d-22 [-1, 512, 28, 28] 2,359,808

ReLU-23 [-1, 512, 28, 28] 0

MaxPool2d-24 [-1, 512, 14, 14] 0

Conv2d-25 [-1, 512, 14, 14] 2,359,808

ReLU-26 [-1, 512, 14, 14] 0

Conv2d-27 [-1, 512, 14, 14] 2,359,808

ReLU-28 [-1, 512, 14, 14] 0

Conv2d-29 [-1, 512, 14, 14] 2,359,808

ReLU-30 [-1, 512, 14, 14] 0

MaxPool2d-31 [-1, 512, 7, 7] 0

Linear-32 [-1, 4096] 102,764,544

ReLU-33 [-1, 4096] 0

Dropout-34 [-1, 4096] 0

Linear-35 [-1, 4096] 16,781,312

ReLU-36 [-1, 4096] 0

Dropout-37 [-1, 4096] 0

Linear-38 [-1, 1000] 4,097,000

================================================================

Total params: 138,357,544

Trainable params: 138,357,544

Non-trainable params: 0

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

Input size (MB): 0.57

Forward/backward pass size (MB): 218.59

Params size (MB): 527.79

Estimated Total Size (MB): 746.96

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