在使用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模型的安装方法为:
pip install torchkeras
调用代码如下:
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()
# 使用torchkeras中的summary函数打印模型结构和参数
print(summary(net, input_shape=(15, )))
效果如下:
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Linear-1 [-1, 20] 320
ReLU-2 [-1, 20] 0
Linear-3 [-1, 1] 21
Sigmoid-4 [-1, 1] 0
================================================================
Total params: 341
Trainable params: 341
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.000057
Forward/backward pass size (MB): 0.000320
Params size (MB): 0.001301
Estimated Total Size (MB): 0.001678
----------------------------------------------------------------
None
可以看到这里模型的层次结构、输出形状和参数基本信息都可以被打印出来,这个模块还是很不错的。此外torchkeras的功能不仅限于此,它是在Pytorch上实现的仿Keras的高层Model结构,可以让我们像使用Keras一样训练Pytorch模型,关于该模块的细节可以参考以下文章:
torchkeras,像Keras一样训练Pytorch模型!