| 659 | |
| 660 | |
| 661 | class StyleEncoder(nn.Module): |
| 662 | def __init__(self, n_downsample, input_dim, dim, style_dim, norm, activ, vae=False): |
| 663 | super(StyleEncoder, self).__init__() |
| 664 | self.vae = vae |
| 665 | self.model = [] |
| 666 | self.model += [Conv2dBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type='reflect')] |
| 667 | for i in range(2): |
| 668 | self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type='reflect')] |
| 669 | dim *= 2 |
| 670 | for i in range(n_downsample - 2): |
| 671 | self.model += [Conv2dBlock(dim, dim, 4, 2, 1, norm=norm, activation=activ, pad_type='reflect')] |
| 672 | self.model += [nn.AdaptiveAvgPool2d(1)] # global average pooling |
| 673 | if self.vae: |
| 674 | self.fc_mean = nn.Linear(dim, style_dim) # , 1, 1, 0) |
| 675 | self.fc_var = nn.Linear(dim, style_dim) # , 1, 1, 0) |
| 676 | else: |
| 677 | self.model += [nn.Conv2d(dim, style_dim, 1, 1, 0)] |
| 678 | |
| 679 | self.model = nn.Sequential(*self.model) |
| 680 | self.output_dim = dim |
| 681 | |
| 682 | def forward(self, x): |
| 683 | if self.vae: |
| 684 | output = self.model(x) |
| 685 | output = output.view(x.size(0), -1) |
| 686 | output_mean = self.fc_mean(output) |
| 687 | output_var = self.fc_var(output) |
| 688 | return output_mean, output_var |
| 689 | else: |
| 690 | return self.model(x).view(x.size(0), -1) |
| 691 | |
| 692 | |
| 693 | class ContentEncoder(nn.Module): |