| 89 | |
| 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): |
| 124 | # alpha should be scalar within [0, 1], and upscale.shape == generated.shape |
| 125 | return torch.tanh(alpha * generated + (1 - alpha) * upscaled) |
| 126 | |
| 127 | def forward(self, x, alpha, steps): |
| 128 | out = self.initial(x) |
| 129 | |
| 130 | if steps == 0: |
| 131 | return self.initial_rgb(out) |
| 132 | |
| 133 | upscaled = 0 |
| 134 | |
| 135 | for step in range(steps): |
| 136 | #每一个convblock之后进行上采样 |
| 137 | upscaled = F.interpolate(out, scale_factor=2, mode="nearest") |
| 138 | #进入下一个convblock |
| 139 | out = self.prog_blocks[step](upscaled) |
| 140 | |
| 141 | """ |
| 142 | # The number of channels in upscale will stay the same, while |
| 143 | # out which has moved through prog_blocks might change. To ensure |
| 144 | # we can convert both to rgb we use different rgb_layers |
| 145 | # (steps-1) and steps for upscaled, out respectively |
| 146 | """ |
| 147 | final_upscaled = self.rgb_layers[steps - 1](upscaled) |
| 148 | final_out = self.rgb_layers[steps](out) |
no outgoing calls
no test coverage detected