本文主要说两个方式
- torchsummary库
- 关于一个可视化的函数,原文链接是:可视化(输出pdf)
由于 我们需要展示,这里我使用CycleGAN的利用残差块的生成器
代码如下:
class ResidualBlock(nn.Module):
def __init__(self, in_features):
super(ResidualBlock, self).__init__()
conv_block = [ nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(1),
nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features) ]
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x):
return x + self.conv_block(x)
class GeneratorResNet(nn.Module):
def __init__(self, in_channels, out_channels, res_blocks ):
super(GeneratorResNet, self).__init__()
#in_channels = args.input_nc
#out_channels = args.output_nc
#res_blocks = args.n_residual_blocks
# Initial convolution block
model = [ nn.ReflectionPad2d(3),
nn.Conv2d(in_channels, 64, 7),
nn.InstanceNorm2d(64),
nn.ReLU(inplace=True) ]
# Downsampling
in_features