安裝torchsummary

pip install torchsummary

例子

netE = Encoder(NUM_CHANNELS=3)
netG = Generator(NUM_CHANNELS=3, z_dim=1024, z_size=1)
netD = Discriminator(NUM_CHANNELS=3, z_dim=1024)
netD.to(device)
netE.to(device)
netG.to(device)
summary(netE, input_size=(3, 128, 128), batch_size=-1)
summary(netG, input_size=(1024, 1, 1), batch_size=-1)
summary(netD, input_size=[(3, 128, 128),(1024, 1, 1)], batch_size=-1)