| 90 | |
| 91 | class Generator(nn.Module): |
| 92 | def __init__(self, z_dim, in_channels, img_channels=3): |
| 93 | super(Generator, self).__init__() |
| 94 | |
| 95 | # initial takes 1x1 -> 4x4 |
| 96 | self.initial = nn.Sequential( |
| 97 | PixelNorm(), |
| 98 | nn.ConvTranspose2d(z_dim, in_channels, kernel_size=(4,4), stride=(1,1) ,padding=(0,0)), |
| 99 | nn.LeakyReLU(0.2), |
| 100 | WSConv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1), |
| 101 | nn.LeakyReLU(0.2), |
| 102 | PixelNorm(), |
| 103 | ) |
| 104 | |
| 105 | self.initial_rgb = WSConv2d( |
| 106 | in_channels, img_channels, kernel_size=1, stride=1, padding=0 |
| 107 | ) |
| 108 | #self.rgb_layers表示每一个convblock之后的将特征图转换为rgb图像 |
| 109 | self.prog_blocks, self.rgb_layers = ( |
| 110 | nn.ModuleList([]), |
| 111 | nn.ModuleList([self.initial_rgb]), |
| 112 | ) |
| 113 | |
| 114 | for i in range(len(factors) - 1): # -1 to prevent index error because of factors[i+1] |
| 115 | conv_in_c = int(in_channels * factors[i]) |
| 116 | conv_out_c = int(in_channels * factors[i + 1]) |
| 117 | self.prog_blocks.append(ConvBlock(conv_in_c, conv_out_c)) |
| 118 | self.rgb_layers.append( |
| 119 | WSConv2d(conv_out_c, img_channels, kernel_size=1, stride=1, padding=0) |
| 120 | ) |
| 121 | |
| 122 | # |
| 123 | def fade_in(self, alpha, upscaled, generated): |