(self, n_downsample, input_dim, dim, style_dim, norm, activ, vae=False)
| 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: |
nothing calls this directly
no test coverage detected