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 分钟前
【EI/Scopus双检索】2025年4月光电信息、传感云、边缘计算、光学成像、物联网、智慧城市、新材料国际学术盛宴来袭!
人工智能·神经网络·物联网·算法·智慧城市·边缘计算
Listennnn9 分钟前
神经网络能不能完全拟合y=x² ???
人工智能·深度学习·神经网络
[shenhonglei]1 小时前
【吉卜力风格Prompt 超好用现成提示词】
人工智能
【云轩】1 小时前
《信号革命:从模拟到数字的通信进化史诗》
人工智能·嵌入式硬件·语音识别
视觉&物联智能1 小时前
【杂谈】-大型语言模型对具身人工智能发展的推动与挑战
人工智能·搜索引擎·语言模型·大模型·llm·具身人工智能
巫山老妖1 小时前
5分钟手把手教你开发一个MCP服务
人工智能
巫山老妖1 小时前
大模型 MCP:开启 AI 与现实世界的无缝交互革命
人工智能
不吃香菜?1 小时前
Opencv之dilib库:表情识别
人工智能·opencv·计算机视觉
郝YH是人间理想2 小时前
OpenCV基础——傅里叶变换、角点检测
开发语言·图像处理·人工智能·python·opencv·计算机视觉
__Benco2 小时前
OpenHarmony子系统开发 - 安全(十)
人工智能·harmonyos