这里记录一个很好使用的工具-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"