fliw2/progan
0
1import torch
2from torch import nn
3from torch.nn import functional as F
4
5from math import sqrt
6
7
8class EqualLR:
9 def __init__(self, name):
10 self.name = name
11
12 def compute_weight(self, module):
13 weight = getattr(module, self.name + '_orig')
14 fan_in = weight.data.size(1) * weight.data[0][0].numel()
15
16 return weight * sqrt(2 / fan_in)
17
18 @staticmethod
19 def apply(module, name):
20 fn = EqualLR(name)
21
22 weight = getattr(module, name)
23 del module._parameters[name]
24 module.register_parameter(name + '_orig', nn.Parameter(weight.data))
25 module.register_forward_pre_hook(fn)
26
27 return fn
28
29 def __call__(self, module, input):
30 weight = self.compute_weight(module)
31 setattr(module, self.name, weight)
32
33
34def equal_lr(module, name='weight'):
35 EqualLR.apply(module, name)
36
37 return module
38
39
40class PixelNorm(nn.Module):
41 def __init__(self):
42 super().__init__()
43
44 def forward(self, input):
45 return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True)
46 + 1e-8)
47
48
49class EqualConv2d(nn.Module):
50 def __init__(self, *args, **kwargs):
51 super().__init__()
52
53 conv = nn.Conv2d(*args, **kwargs)
54 conv.weight.data.normal_()
55 conv.bias.data.zero_()
56 self.conv = equal_lr(conv)
57
58 def forward(self, input):
59 return self.conv(input)
60
61
62class EqualConvTranspose2d(nn.Module):
63 ### additional module for OOGAN usage
64 def __init__(self, *args, **kwargs):
65 super().__init__()
66
67 conv = nn.ConvTranspose2d(*args, **kwargs)
68 conv.weight.data.normal_()
69 conv.bias.data.zero_()
70 self.conv = equal_lr(conv)
71
72 def forward(self, input):
73 return self.conv(input)
74
75class EqualLinear(nn.Module):
76 def __init__(self, in_dim, out_dim):
77 super().__init__()
78
79 linear = nn.Linear(in_dim, out_dim)
80 linear.weight.data.normal_()
81 linear.bias.data.zero_()
82
83 self.linear = equal_lr(linear)
84
85 def forward(self, input):
86 return self.linear(input)
87
88
89class ConvBlock(nn.Module):
90 def __init__(self, in_channel, out_channel, kernel_size, padding, kernel_size2=None, padding2=None, pixel_norm=True):
91 super().__init__()
92
93 pad1 = padding
94 pad2 = padding
95 if padding2 is not None:
96 pad2 = padding2
97
98 kernel1 = kernel_size
99 kernel2 = kernel_size
100 if kernel_size2 is not None:
101 kernel2 = kernel_size2
102
103 convs = [EqualConv2d(in_channel, out_channel, kernel1, padding=pad1)]
104 if pixel_norm:
105 convs.append(PixelNorm())
106 convs.append(nn.LeakyReLU(0.1))
107 convs.append(EqualConv2d(out_channel, out_channel, kernel2, padding=pad2))
108 if pixel_norm:
109 convs.append(PixelNorm())
110 convs.append(nn.LeakyReLU(0.1))
111
112 self.conv = nn.Sequential(*convs)
113
114 def forward(self, input):
115 out = self.conv(input)
116 return out
117
118
119def upscale(feat):
120 return F.interpolate(feat, scale_factor=2, mode='bilinear', align_corners=False)
121
122class Generator(nn.Module):
123 def __init__(self, input_code_dim=128, in_channel=128, pixel_norm=True, tanh=True):
124 super().__init__()
125 self.input_dim = input_code_dim
126 self.tanh = tanh
127 self.input_layer = nn.Sequential(
128 EqualConvTranspose2d(input_code_dim, in_channel, 4, 1, 0),
129 PixelNorm(),
130 nn.LeakyReLU(0.1))
131
132 self.progression_4 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
133 self.progression_8 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
134 self.progression_16 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
135 self.progression_32 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
136 self.progression_64 = ConvBlock(in_channel, in_channel//2, 3, 1, pixel_norm=pixel_norm)
137 self.progression_128 = ConvBlock(in_channel//2, in_channel//4, 3, 1, pixel_norm=pixel_norm)
138 self.progression_256 = ConvBlock(in_channel//4, in_channel//4, 3, 1, pixel_norm=pixel_norm)
139
140 self.to_rgb_8 = EqualConv2d(in_channel, 3, 1)
141 self.to_rgb_16 = EqualConv2d(in_channel, 3, 1)
142 self.to_rgb_32 = EqualConv2d(in_channel, 3, 1)
143 self.to_rgb_64 = EqualConv2d(in_channel//2, 3, 1)
144 self.to_rgb_128 = EqualConv2d(in_channel//4, 3, 1)
145 self.to_rgb_256 = EqualConv2d(in_channel//4, 3, 1)
146
147 self.max_step = 6
148
149 def progress(self, feat, module):
150 out = F.interpolate(feat, scale_factor=2, mode='bilinear', align_corners=False)
151 out = module(out)
152 return out
153
154 def output(self, feat1, feat2, module1, module2, alpha):
155 if 0 <= alpha < 1:
156 skip_rgb = upscale(module1(feat1))
157 out = (1-alpha)*skip_rgb + alpha*module2(feat2)
158 else:
159 out = module2(feat2)
160 if self.tanh:
161 return torch.tanh(out)
162 return out
163
164 def forward(self, input, step=0, alpha=-1):
165 if step > self.max_step:
166 step = self.max_step
167
168 out_4 = self.input_layer(input.view(-1, self.input_dim, 1, 1))
169 out_4 = self.progression_4(out_4)
170 out_8 = self.progress(out_4, self.progression_8)
171 if step==1:
172 if self.tanh:
173 return torch.tanh(self.to_rgb_8(out_8))
174 return self.to_rgb_8(out_8)
175
176 out_16 = self.progress(out_8, self.progression_16)
177 if step==2:
178 return self.output( out_8, out_16, self.to_rgb_8, self.to_rgb_16, alpha )
179
180 out_32 = self.progress(out_16, self.progression_32)
181 if step==3:
182 return self.output( out_16, out_32, self.to_rgb_16, self.to_rgb_32, alpha )
183
184 out_64 = self.progress(out_32, self.progression_64)
185 if step==4:
186 return self.output( out_32, out_64, self.to_rgb_32, self.to_rgb_64, alpha )
187
188 out_128 = self.progress(out_64, self.progression_128)
189 if step==5:
190 return self.output( out_64, out_128, self.to_rgb_64, self.to_rgb_128, alpha )
191
192 out_256 = self.progress(out_128, self.progression_256)
193 if step==6:
194 return self.output( out_128, out_256, self.to_rgb_128, self.to_rgb_256, alpha )
195
196
197class Discriminator(nn.Module):
198 def __init__(self, feat_dim=128):
199 super().__init__()
200
201 self.progression = nn.ModuleList([ConvBlock(feat_dim//4, feat_dim//4, 3, 1),
202 ConvBlock(feat_dim//4, feat_dim//2, 3, 1),
203 ConvBlock(feat_dim//2, feat_dim, 3, 1),
204 ConvBlock(feat_dim, feat_dim, 3, 1),
205 ConvBlock(feat_dim, feat_dim, 3, 1),
206 ConvBlock(feat_dim, feat_dim, 3, 1),
207 ConvBlock(feat_dim+1, feat_dim, 3, 1, 4, 0)])
208
209 self.from_rgb = nn.ModuleList([EqualConv2d(3, feat_dim//4, 1),
210 EqualConv2d(3, feat_dim//4, 1),
211 EqualConv2d(3, feat_dim//2, 1),
212 EqualConv2d(3, feat_dim, 1),
213 EqualConv2d(3, feat_dim, 1),
214 EqualConv2d(3, feat_dim, 1),
215 EqualConv2d(3, feat_dim, 1)])
216
217 self.n_layer = len(self.progression)
218
219 self.linear = EqualLinear(feat_dim, 1)
220
221 def forward(self, input, step=0, alpha=-1):
222 for i in range(step, -1, -1):
223 index = self.n_layer - i - 1
224
225 if i == step:
226 out = self.from_rgb[index](input)
227
228 if i == 0:
229 out_std = torch.sqrt(out.var(0, unbiased=False) + 1e-8)
230 mean_std = out_std.mean()
231 mean_std = mean_std.expand(out.size(0), 1, 4, 4)
232 out = torch.cat([out, mean_std], 1)
233
234 out = self.progression[index](out)
235
236 if i > 0:
237 # out = F.avg_pool2d(out, 2)
238 out = F.interpolate(out, scale_factor=0.5, mode='bilinear', align_corners=False)
239
240 if i == step and 0 <= alpha < 1:
241 # skip_rgb = F.avg_pool2d(input, 2)
242 skip_rgb = F.interpolate(input, scale_factor=0.5, mode='bilinear', align_corners=False)
243 skip_rgb = self.from_rgb[index + 1](skip_rgb)
244 out = (1 - alpha) * skip_rgb + alpha * out
245
246 out = out.squeeze(2).squeeze(2)
247 # print(input.size(), out.size(), step)
248 out = self.linear(out)
249
250 return out