目录
一、旧 summary
pip install torchsummary
使用方法:(但是加入LSTM后会报错)
from torchsummary import summary
net = Resnet(basic_block, [2, 2, 2, 2], 16).cuda()
summary(net, (3, 9, 9))
二、新 summary
pip install torchinfo
使用方法:(加入LSTM后不会报错)
from torchinfo import summary
model = Resnet(basic_block, [2, 2, 2, 2], 16).cuda()
batch_size = 16
summary(model, input_size=(batch_size, 1, 28, 28))