pytorch 模型可视化

用torchsummary库实现。summary(model, input_size=(channel, height, width)

from torchsummary import summary
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = UNet()
model = model.to(device)
summary(model, input_size=(1, 480, 480))

输出如下：