pytorch基础学习-torchsummary计算参数量

这里记录一个很好使用的工具-torchsummary，可以查看每一层网络的结构及其大小、参数量，最后可以计算整个网络的大小、参数量。

1、安装

pip install torchsummary

2、测试

python模式下输入import torchsummary

3、使用

import torchsummary

net = xxx # 自己的网络
print(torchsummary.summary(net, (3, 256, 256), device="cuda"))

• torchsummary.summary(model, input_size, batch_size=-1, device="cuda")
• model：自己的网络模型
• input_size：输入尺寸，形状为 C，H ，W，分别为通道数、高、宽
• batch_size：batch_size，默认为 -1，在展示模型每层输出的形状时显示的 batch_size
• device：选择"cuda"或"cpu"

4、结果展示