| 62 | |
| 63 | |
| 64 | class ADAIN_Encoder(nn.Module): |
| 65 | def __init__(self, encoder, gpu_ids=[]): |
| 66 | super(ADAIN_Encoder, self).__init__() |
| 67 | enc_layers = list(encoder.children()) |
| 68 | self.enc_1 = nn.Sequential(*enc_layers[:4]) # input -> relu1_1 64 |
| 69 | self.enc_2 = nn.Sequential(*enc_layers[4:11]) # relu1_1 -> relu2_1 128 |
| 70 | self.enc_3 = nn.Sequential(*enc_layers[11:18]) # relu2_1 -> relu3_1 256 |
| 71 | self.enc_4 = nn.Sequential(*enc_layers[18:31]) # relu3_1 -> relu4_1 512 |
| 72 | |
| 73 | self.mse_loss = nn.MSELoss() |
| 74 | |
| 75 | # fix the encoder |
| 76 | for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4']: |
| 77 | for param in getattr(self, name).parameters(): |
| 78 | param.requires_grad = False |
| 79 | |
| 80 | # extract relu1_1, relu2_1, relu3_1, relu4_1 from input image |
| 81 | def encode_with_intermediate(self, input): |
| 82 | results = [input] |
| 83 | for i in range(4): |
| 84 | func = getattr(self, 'enc_{:d}'.format(i + 1)) |
| 85 | results.append(func(results[-1])) |
| 86 | return results[1:] |
| 87 | |
| 88 | def calc_mean_std(self, feat, eps=1e-5): |
| 89 | # eps is a small value added to the variance to avoid divide-by-zero. |
| 90 | size = feat.size() |
| 91 | assert (len(size) == 4) |
| 92 | N, C = size[:2] |
| 93 | feat_var = feat.view(N, C, -1).var(dim=2) + eps |
| 94 | feat_std = feat_var.sqrt().view(N, C, 1, 1) |
| 95 | feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) |
| 96 | return feat_mean, feat_std |
| 97 | |
| 98 | def adain(self, content_feat, style_feat): |
| 99 | assert (content_feat.size()[:2] == style_feat.size()[:2]) |
| 100 | size = content_feat.size() |
| 101 | style_mean, style_std = self.calc_mean_std(style_feat) |
| 102 | content_mean, content_std = self.calc_mean_std(content_feat) |
| 103 | |
| 104 | normalized_feat = (content_feat - content_mean.expand( |
| 105 | size)) / content_std.expand(size) |
| 106 | return normalized_feat * style_std.expand(size) + style_mean.expand(size) |
| 107 | |
| 108 | def forward(self, content, style, encoded_only = False): |
| 109 | style_feats = self.encode_with_intermediate(style) |
| 110 | content_feats = self.encode_with_intermediate(content) |
| 111 | if encoded_only: |
| 112 | return content_feats[-1], style_feats[-1] |
| 113 | else: |
| 114 | adain_feat = self.adain(content_feats[-1], style_feats[-1]) |
| 115 | return adain_feat |
| 116 | |
| 117 | class Decoder(nn.Module): |
| 118 | def __init__(self, gpu_ids=[]): |
nothing calls this directly
no outgoing calls
no test coverage detected