| 691 | |
| 692 | |
| 693 | class ContentEncoder(nn.Module): |
| 694 | def __init__(self, n_downsample, n_res, input_dim, dim, norm, activ, pad_type='zero'): |
| 695 | super(ContentEncoder, self).__init__() |
| 696 | self.model = [] |
| 697 | self.model += [Conv2dBlock(input_dim, dim, 7, 1, 3, norm=norm, activation=activ, pad_type='reflect')] |
| 698 | # downsampling blocks |
| 699 | for i in range(n_downsample): |
| 700 | self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1, norm=norm, activation=activ, pad_type='reflect')] |
| 701 | dim *= 2 |
| 702 | # residual blocks |
| 703 | self.model += [ResBlocks(n_res, dim, norm=norm, activation=activ, pad_type=pad_type)] |
| 704 | self.model = nn.Sequential(*self.model) |
| 705 | self.output_dim = dim |
| 706 | |
| 707 | def forward(self, x, nce_layers=[], encode_only=False): |
| 708 | if len(nce_layers) > 0: |
| 709 | feat = x |
| 710 | feats = [] |
| 711 | for layer_id, layer in enumerate(self.model): |
| 712 | feat = layer(feat) |
| 713 | if layer_id in nce_layers: |
| 714 | feats.append(feat) |
| 715 | if layer_id == nce_layers[-1] and encode_only: |
| 716 | return None, feats |
| 717 | return feat, feats |
| 718 | else: |
| 719 | return self.model(x), None |
| 720 | |
| 721 | class Decoder_all(nn.Module): |
| 722 | def __init__(self, n_upsample, n_res, dim, output_dim, norm='batch', activ='relu', pad_type='zero', nz=0): |