【676】PyTorch —— 网络结构显示

Pytorch 查看模型网络结构

　　有两种方式，不过效果都没有 keras 自带函数效果好，凑活用吧!!!

方法一：直接打印

print(model)

　　效果：

方法二：通过 torchsummary

from torchsummary import summary 

model = UNet(3, 2)
summary(model, (3, 512, 512))

　　效果：