Pytorch如何打印与Keras的model.summary()类似的输出

1 Keras的model.summary()

2 Pytorch实现

2.1 安装torchsummary包

powershell 复制代码
pip install torchsummary

2.2 代码

python 复制代码
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))

2.3 运行结果

参考材料:

相关推荐
~央千澈~1 小时前
对鸿蒙 Next 系统“成熟论”的深度剖析-优雅草卓伊凡
人工智能
Donvink1 小时前
【视频生成模型】通义万相Wan2.1模型本地部署和LoRA微调
人工智能·深度学习·aigc·音视频
訾博ZiBo1 小时前
AI日报 - 2025年04月29日
人工智能
爱喝奶茶的企鹅1 小时前
Ethan独立开发产品日报 | 2025-04-27
人工智能·程序员·开源
极小狐1 小时前
如何对极狐GitLab 议题进行过滤和排序?
人工智能·git·机器学习·gitlab
望获linux2 小时前
智能清洁机器人中的实时操作系统应用研究
大数据·linux·服务器·人工智能·机器人·操作系统
古月居GYH2 小时前
用高斯溅射技术跨越机器人模拟与现实的鸿沟:SplatSim 框架解析
人工智能·机器人
槑槑紫2 小时前
pytorch(gpu版本安装)
人工智能·pytorch·python
飞哥数智坊2 小时前
五一必备:手把手教你“即梦”APP轻松生成精美海报
人工智能