- 安装第三方库torchsummary
pip install torchsummary
- 导入库
from torchsummary import summary
- 获取pytorch模型参数情况
summary(model, input_size=(channels, height, width), batch_size=numbers)
其中,channels是输入张量的channels数,height是输入张量的高,width是输入张量的宽,numbers是模型的batchsize数。