| 157 | |
| 158 | |
| 159 | class UpSampleResNetSkip(BaseModule): |
| 160 | def __init__(self, nc_in, nc_out, nf, use_bias, norm, conv_by, conv_type, |
| 161 | use_skip_connection=False): |
| 162 | super().__init__(conv_type) |
| 163 | assert False, "UpSampleResNetSkip has worse performance than UpSampleResNetSkipGated" |
| 164 | # Upsample 1 |
| 165 | self.conv_c2 = self.ConvBlock(512, 256, kernel_size=(3, 1, 1), stride=1, |
| 166 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 167 | self.conv_c4 = self.ConvBlock(1024, 256, kernel_size=(3, 1, 1), stride=1, |
| 168 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 169 | |
| 170 | self.deconv1 = self.DeconvBlock(2048, 256, kernel_size=(3, 1, 1), stride=1, |
| 171 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 172 | self.conv9 = self.ConvBlock( |
| 173 | 256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 174 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 175 | |
| 176 | # Upsample 2 |
| 177 | self.deconv2 = self.DeconvBlock( |
| 178 | 256, |
| 179 | 256, kernel_size=(3, 1, 1), stride=1, |
| 180 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 181 | self.conv10 = self.ConvBlock( |
| 182 | 256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1, |
| 183 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 184 | |
| 185 | # Upsample 3 |
| 186 | self.deconv3 = self.DeconvBlock( |
| 187 | 256, |
| 188 | 256, kernel_size=(3, 1, 1), stride=1, |
| 189 | bias=use_bias, norm=norm, conv_by=conv_by) |
| 190 | self.conv11 = self.ConvBlock( |
| 191 | 256, nc_out, kernel_size=(3, 3, 3), stride=(1, 1, 1), |
| 192 | padding=1, bias=use_bias, norm=None, activation=None, conv_by=conv_by) |
| 193 | |
| 194 | for name, module in self.named_modules(): |
| 195 | if isinstance(module, (VanillaConv)): |
| 196 | nn.init.kaiming_uniform_(module.featureConv.layer.weight, a=1) |
| 197 | nn.init.constant_(module.featureConv.layer.bias, 0) |
| 198 | elif isinstance(module, (VanillaDeconv)): |
| 199 | nn.init.kaiming_uniform_(module.conv.featureConv.layer.weight, a=1) |
| 200 | nn.init.constant_(module.conv.featureConv.layer.bias, 0) |
| 201 | |
| 202 | def forward(self, inp): |
| 203 | c1, c2, c4, c8 = inp |
| 204 | |
| 205 | c2 = self.conv_c2(c2) |
| 206 | c4 = self.conv_c4(c4) |
| 207 | |
| 208 | c4 = interpolate(c4, scale_factor=2, mode="nearest") |
| 209 | c2 = interpolate(c2, scale_factor=4, mode="nearest") |
| 210 | c1 = interpolate(c1, scale_factor=4, mode="nearest") |
| 211 | |
| 212 | d1 = self.deconv1(c8) |
| 213 | d1 = d1 + c4 |
| 214 | c9 = self.conv9(d1) |
| 215 | d2 = self.deconv2(c9) |
| 216 | d2 = d2 + c2 |