使用torchkeras打印Pytorch模型结构和基本参数信息

最新推荐文章于 2025-11-17 15:43:42 发布

原创

最新推荐文章于 2025-11-17 15:43:42 发布 · 5.3k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #模型结构 #torchkeras #python #深度学习

本文介绍如何使用torchkeras模块的summary函数查看Pytorch构建的神经网络模型结构，包括层次结构、输出形状和参数基本信息，使模型结构显示更加详细直观。

在使用Pytorch构建神经网络模型后，我们需要看一下自己写的模型的网络结构，此时可以使用torchkeras模块中的summary函数实现该功能。以多层感知机为例，首先我们构建网络结构并打印该模型的初步信息，代码如下：

import torch
from torch import nn
from torchkeras import summary


def create_net():
    net = nn.Sequential()
    net.add_module('linear1', nn.Linear(15, 20))
    net.add_module('relu1', nn.ReLU())
    net.add_module('linear2', nn.Linear(20, 1))
    net.add_module('sigmoid', nn.Sigmoid())
    return net

# 创建模型
net = create_net()
# 打印模型的基本信息
print(net)

效果如下：

Sequential(
  (linear1): Linear(in_features=15, out_features=20, bias=True)
  (relu1): ReLU()
  (linear3): Linear(in_features=20, out_features=1, bias=True)
  (sigmoid): Sigmoid()
)

可以看到这种方式输出的模型结构内容不够详细而且形式不够直观，我们想实现像keras打印模型结构的那种效果，可以使用torchkeras模型的summary函数。

torchkeras模型的安装方法为：