深度学习各个部分数据大小抓取方法
2021-01-07 本文已影响0人
CPinging
一、使用Summary方案
效果如图：
image.png
第一列为Layer名称，第二列为输出数据的维度，第三列为weight的数量。
具体方法：
-
1
pip3 install torchsummary； -
2
from torchsummary import summary； -
3 然后在训练后添加summary这句话，后面传入net定义的模型，并传入dataset的维度即可打印出来；
...
net.train()
summary(net, (3, 32, 32))
...
