| 123 | |
| 124 | #对应着Synthesis network的每一个block |
| 125 | class GenBlock(torch.nn.Module): |
| 126 | def __init__(self,in_channels,out_channels,w_dim): |
| 127 | super(GenBlock, self).__init__() |
| 128 | self.conv1 = WSConv2d(in_channels,out_channels) |
| 129 | self.conv2 = WSConv2d(out_channels,out_channels) |
| 130 | self.leaky = torch.nn.LeakyReLU(negative_slope=0.2,inplace=True) |
| 131 | self.inject_noise1 = InjectNoise(out_channels) |
| 132 | self.inject_noise2 = InjectNoise(out_channels) |
| 133 | self.adain1 = AdaIN(out_channels,w_dim) |
| 134 | self.adain2 = AdaIN(out_channels,w_dim) |
| 135 | |
| 136 | def forward(self,x,w): |
| 137 | x = self.adain1(self.leaky(self.inject_noise1(self.conv1(x))),w) |
| 138 | x = self.adain2(self.leaky(self.inject_noise2(self.conv2(x))),w) |
| 139 | return x |
| 140 | |
| 141 | class ConvBlock(torch.nn.Module): |
| 142 | def __init__(self,in_channels,out_channels): |