pytorch查看网络结构图

方法1:

https://github.com/Lyken17/pytorch-OpCounter

from torchvision.models import resnet50	
from thop import profile	
	
model = resnet50()	
input = torch.randn(1, 3, 224, 224)	
flops, params = profile(model, inputs=(input, ))

方法2:

import torch

import torchvision.models as models

from torchsummary import summary

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

model = models.resnet18().to(device)

summary(model, (3,256,256))

 

 

 

善始者众,善终者寡

 

  • 3
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
PyTorch自定义网络模型结构可以通过使用PyTorch内置的`torchsummary`模块来生成。这个模块可以帮助我们快速地展示模型的参数数量、每一层的输出形状等重要信息。 以下是一个简单的示例,展示了如何使用`torchsummary`模块来生成自定义网络模型的结构: ``` python import torch import torch.nn as nn from torchsummary import summary class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) self.fc1 = nn.Linear(64 * 8 * 8, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = nn.functional.relu(x) x = self.pool(x) x = self.conv2(x) x = nn.functional.relu(x) x = self.pool(x) x = x.view(-1, 64 * 8 * 8) x = self.fc1(x) x = nn.functional.relu(x) x = self.fc2(x) return x model = MyModel() summary(model, input_size=(3, 32, 32)) ``` 运行以上代码,就可以得到如下的输出: ``` ---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 32, 32, 32] 896 MaxPool2d-2 [-1, 32, 16, 16] 0 Conv2d-3 [-1, 64, 16, 16] 18,496 MaxPool2d-4 [-1, 64, 8, 8] 0 Linear-5 [-1, 128] 524,416 Linear-6 [-1, 10] 1,290 ================================================================ Total params: 545,098 Trainable params: 545,098 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 0.01 Forward/backward pass size (MB): 0.75 Params size (MB): 2.08 Estimated Total Size (MB): 2.85 ---------------------------------------------------------------- ``` 可以看到,`summary`函数生成了一个包含每一层输出形状、参数数量等信息的表格,以及估计的模型大小。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值