ICML2022/resefa
4
1# python3.72"""Contains the implementation of encoder used in GH-Feat (including IDInvert).3 4ResNet is used as the backbone.5 6GH-Feat paper: https://arxiv.org/pdf/2007.10379.pdf7IDInvert paper: https://arxiv.org/pdf/2004.00049.pdf8 9NOTE: Please use `latent_num` and `num_latents_per_head` to control the10inversion space, such as Y-space used in GH-Feat and W-space used in IDInvert.11In addition, IDInvert sets `use_fpn` and `use_sam` as `False` by default.12"""13 14import numpy as np15 16import torch17import torch.nn as nn18import torch.nn.functional as F19import torch.distributed as dist20 21__all__ = ['GHFeatEncoder']22 23# Resolutions allowed.24_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]25 26# pylint: disable=missing-function-docstring27 28class BasicBlock(nn.Module):29 """Implementation of ResNet BasicBlock."""30 31 expansion = 132 33 def __init__(self,34 inplanes,35 planes,36 base_width=64,37 stride=1,38 groups=1,39 dilation=1,40 norm_layer=None,41 downsample=None):42 super().__init__()43 if base_width != 64:44 raise ValueError(f'BasicBlock of ResNet only supports '45 f'`base_width=64`, but {base_width} received!')46 if stride not in [1, 2]:47 raise ValueError(f'BasicBlock of ResNet only supports `stride=1` '48 f'and `stride=2`, but {stride} received!')49 if groups != 1:50 raise ValueError(f'BasicBlock of ResNet only supports `groups=1`, '51 f'but {groups} received!')52 if dilation != 1:53 raise ValueError(f'BasicBlock of ResNet only supports '54 f'`dilation=1`, but {dilation} received!')55 assert self.expansion == 156 57 self.stride = stride58 if norm_layer is None:59 norm_layer = nn.BatchNorm2d60 self.conv1 = nn.Conv2d(in_channels=inplanes,61 out_channels=planes,62 kernel_size=3,63 stride=stride,64 padding=1,65 groups=1,66 dilation=1,67 bias=False)68 self.bn1 = norm_layer(planes)69 self.relu = nn.ReLU(inplace=True)70 self.conv2 = nn.Conv2d(in_channels=planes,71 out_channels=planes,72 kernel_size=3,73 stride=1,74 padding=1,75 groups=1,76 dilation=1,77 bias=False)78 self.bn2 = norm_layer(planes)79 self.downsample = downsample80 81 def forward(self, x):82 identity = self.downsample(x) if self.downsample is not None else x83 84 out = self.conv1(x)85 out = self.bn1(out)86 out = self.relu(out)87 88 out = self.conv2(out)89 out = self.bn2(out)90 out = self.relu(out + identity)91 92 return out93 94 95class Bottleneck(nn.Module):96 """Implementation of ResNet Bottleneck."""97 98 expansion = 499 100 def __init__(self,101 inplanes,102 planes,103 base_width=64,104 stride=1,105 groups=1,106 dilation=1,107 norm_layer=None,108 downsample=None):109 super().__init__()110 if stride not in [1, 2]:111 raise ValueError(f'Bottleneck of ResNet only supports `stride=1` '112 f'and `stride=2`, but {stride} received!')113 114 width = int(planes * (base_width / 64)) * groups115 self.stride = stride116 if norm_layer is None:117 norm_layer = nn.BatchNorm2d118 self.conv1 = nn.Conv2d(in_channels=inplanes,119 out_channels=width,120 kernel_size=1,121 stride=1,122 padding=0,123 dilation=1,124 groups=1,125 bias=False)126 self.bn1 = norm_layer(width)127 self.conv2 = nn.Conv2d(in_channels=width,128 out_channels=width,129 kernel_size=3,130 stride=stride,131 padding=dilation,132 groups=groups,133 dilation=dilation,134 bias=False)135 self.bn2 = norm_layer(width)136 self.conv3 = nn.Conv2d(in_channels=width,137 out_channels=planes * self.expansion,138 kernel_size=1,139 stride=1,140 padding=0,141 dilation=1,142 groups=1,143 bias=False)144 self.bn3 = norm_layer(planes * self.expansion)145 self.relu = nn.ReLU(inplace=True)146 self.downsample = downsample147 148 def forward(self, x):149 identity = self.downsample(x) if self.downsample is not None else x150 151 out = self.conv1(x)152 out = self.bn1(out)153 out = self.relu(out)154 155 out = self.conv2(out)156 out = self.bn2(out)157 out = self.relu(out)158 159 out = self.conv3(out)160 out = self.bn3(out)161 out = self.relu(out + identity)162 163 return out164 165 166class GHFeatEncoder(nn.Module):167 """Define the ResNet-based encoder network for GAN inversion.168 169 On top of the backbone, there are several task-heads to produce inverted170 codes. Please use `latent_dim` and `num_latents_per_head` to define the171 structure. For example, `latent_dim = [512] * 14` and172 `num_latents_per_head = [4, 4, 6]` can be used for StyleGAN inversion with173 14-layer latent codes, where 3 task heads (corresponding to 4, 4, 6 layers,174 respectively) are used.175 176 Settings for the encoder network:177 178 (1) resolution: The resolution of the output image.179 (2) latent_dim: Dimension of the latent space. A number (one code will be180 produced), or a list of numbers regarding layer-wise latent codes.181 (3) num_latents_per_head: Number of latents that is produced by each head.182 (4) image_channels: Number of channels of the output image. (default: 3)183 (5) final_res: Final resolution of the convolutional layers. (default: 4)184 185 ResNet-related settings:186 187 (1) network_depth: Depth of the network, like 18 for ResNet18. (default: 18)188 (2) inplanes: Number of channels of the first convolutional layer.189 (default: 64)190 (3) groups: Groups of the convolution, used in ResNet. (default: 1)191 (4) width_per_group: Number of channels per group, used in ResNet.192 (default: 64)193 (5) replace_stride_with_dilation: Whether to replace stride with dilation,194 used in ResNet. (default: None)195 (6) norm_layer: Normalization layer used in the encoder. If set as `None`,196 `nn.BatchNorm2d` will be used. Also, please NOTE that when using batch197 normalization, the batch size is required to be larger than one for198 training. (default: nn.BatchNorm2d)199 (7) max_channels: Maximum number of channels in each layer. (default: 512)200 201 Task-head related settings:202 203 (1) use_fpn: Whether to use Feature Pyramid Network (FPN) before outputting204 the latent code. (default: True)205 (2) fpn_channels: Number of channels used in FPN. (default: 512)206 (3) use_sam: Whether to use Spatial Alignment Module (SAM) before outputting207 the latent code. (default: True)208 (4) sam_channels: Number of channels used in SAM. (default: 512)209 """210 211 arch_settings = {212 18: (BasicBlock, [2, 2, 2, 2]),213 34: (BasicBlock, [3, 4, 6, 3]),214 50: (Bottleneck, [3, 4, 6, 3]),215 101: (Bottleneck, [3, 4, 23, 3]),216 152: (Bottleneck, [3, 8, 36, 3])217 }218 219 def __init__(self,220 resolution,221 latent_dim,222 num_latents_per_head,223 image_channels=3,224 final_res=4,225 network_depth=18,226 inplanes=64,227 groups=1,228 width_per_group=64,229 replace_stride_with_dilation=None,230 norm_layer=nn.BatchNorm2d,231 max_channels=512,232 use_fpn=True,233 fpn_channels=512,234 use_sam=True,235 sam_channels=512):236 super().__init__()237 238 if resolution not in _RESOLUTIONS_ALLOWED:239 raise ValueError(f'Invalid resolution: `{resolution}`!\n'240 f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')241 if network_depth not in self.arch_settings:242 raise ValueError(f'Invalid network depth: `{network_depth}`!\n'243 f'Options allowed: '244 f'{list(self.arch_settings.keys())}.')245 if isinstance(latent_dim, int):246 latent_dim = [latent_dim]247 assert isinstance(latent_dim, (list, tuple))248 assert isinstance(num_latents_per_head, (list, tuple))249 assert sum(num_latents_per_head) == len(latent_dim)250 251 self.resolution = resolution252 self.latent_dim = latent_dim253 self.num_latents_per_head = num_latents_per_head254 self.num_heads = len(self.num_latents_per_head)255 self.image_channels = image_channels256 self.final_res = final_res257 self.inplanes = inplanes258 self.network_depth = network_depth259 self.groups = groups260 self.dilation = 1261 self.base_width = width_per_group262 self.replace_stride_with_dilation = replace_stride_with_dilation263 if norm_layer is None:264 norm_layer = nn.BatchNorm2d265 if norm_layer == nn.BatchNorm2d and dist.is_initialized():266 norm_layer = nn.SyncBatchNorm267 self.norm_layer = norm_layer268 self.max_channels = max_channels269 self.use_fpn = use_fpn270 self.fpn_channels = fpn_channels271 self.use_sam = use_sam272 self.sam_channels = sam_channels273 274 block_fn, num_blocks_per_stage = self.arch_settings[network_depth]275 276 self.num_stages = int(np.log2(resolution // final_res)) - 1277 # Add one block for additional stages.278 for i in range(len(num_blocks_per_stage), self.num_stages):279 num_blocks_per_stage.append(1)280 if replace_stride_with_dilation is None:281 replace_stride_with_dilation = [False] * self.num_stages282 283 # Backbone.284 self.conv1 = nn.Conv2d(in_channels=self.image_channels,285 out_channels=self.inplanes,286 kernel_size=7,287 stride=2,288 padding=3,289 bias=False)290 self.bn1 = norm_layer(self.inplanes)291 self.relu = nn.ReLU(inplace=True)292 self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)293 294 self.stage_channels = [self.inplanes]295 self.stages = nn.ModuleList()296 for i in range(self.num_stages):297 inplanes = self.inplanes if i == 0 else planes * block_fn.expansion298 planes = min(self.max_channels, self.inplanes * (2 ** i))299 num_blocks = num_blocks_per_stage[i]300 stride = 1 if i == 0 else 2301 dilate = replace_stride_with_dilation[i]302 self.stages.append(self._make_stage(block_fn=block_fn,303 inplanes=inplanes,304 planes=planes,305 num_blocks=num_blocks,306 stride=stride,307 dilate=dilate))308 self.stage_channels.append(planes * block_fn.expansion)309 310 if self.num_heads > len(self.stage_channels):311 raise ValueError('Number of task heads is larger than number of '312 'stages! Please reduce the number of heads.')313 314 # Task-head.315 if self.num_heads == 1:316 self.use_fpn = False317 self.use_sam = False318 319 if self.use_fpn:320 fpn_pyramid_channels = self.stage_channels[-self.num_heads:]321 self.fpn = FPN(pyramid_channels=fpn_pyramid_channels,322 out_channels=self.fpn_channels)323 if self.use_sam:324 if self.use_fpn:325 sam_pyramid_channels = [self.fpn_channels] * self.num_heads326 else:327 sam_pyramid_channels = self.stage_channels[-self.num_heads:]328 self.sam = SAM(pyramid_channels=sam_pyramid_channels,329 out_channels=self.sam_channels)330 331 self.heads = nn.ModuleList()332 for head_idx in range(self.num_heads):333 # Parse in_channels.334 if self.use_sam:335 in_channels = self.sam_channels336 elif self.use_fpn:337 in_channels = self.fpn_channels338 else:339 in_channels = self.stage_channels[head_idx - self.num_heads]340 in_channels = in_channels * final_res * final_res341 342 # Parse out_channels.343 start_latent_idx = sum(self.num_latents_per_head[:head_idx])344 end_latent_idx = sum(self.num_latents_per_head[:head_idx + 1])345 out_channels = sum(self.latent_dim[start_latent_idx:end_latent_idx])346 347 self.heads.append(CodeHead(in_channels=in_channels,348 out_channels=out_channels,349 norm_layer=self.norm_layer))350 351 def _make_stage(self,352 block_fn,353 inplanes,354 planes,355 num_blocks,356 stride,357 dilate):358 norm_layer = self.norm_layer359 downsample = None360 previous_dilation = self.dilation361 if dilate:362 self.dilation *= stride363 stride = 1364 if stride != 1 or inplanes != planes * block_fn.expansion:365 downsample = nn.Sequential(366 nn.Conv2d(in_channels=inplanes,367 out_channels=planes * block_fn.expansion,368 kernel_size=1,369 stride=stride,370 padding=0,371 dilation=1,372 groups=1,373 bias=False),374 norm_layer(planes * block_fn.expansion),375 )376 377 blocks = []378 blocks.append(block_fn(inplanes=inplanes,379 planes=planes,380 base_width=self.base_width,381 stride=stride,382 groups=self.groups,383 dilation=previous_dilation,384 norm_layer=norm_layer,385 downsample=downsample))386 for _ in range(1, num_blocks):387 blocks.append(block_fn(inplanes=planes * block_fn.expansion,388 planes=planes,389 base_width=self.base_width,390 stride=1,391 groups=self.groups,392 dilation=self.dilation,393 norm_layer=norm_layer,394 downsample=None))395 396 return nn.Sequential(*blocks)397 398 def forward(self, x):399 x = self.conv1(x)400 x = self.bn1(x)401 x = self.relu(x)402 x = self.maxpool(x)403 404 features = [x]405 for i in range(self.num_stages):406 x = self.stages[i](x)407 features.append(x)408 features = features[-self.num_heads:]409 410 if self.use_fpn:411 features = self.fpn(features)412 if self.use_sam:413 features = self.sam(features)414 else:415 final_size = features[-1].shape[2:]416 for i in range(self.num_heads - 1):417 features[i] = F.adaptive_avg_pool2d(features[i], final_size)418 419 outputs = []420 for head_idx in range(self.num_heads):421 codes = self.heads[head_idx](features[head_idx])422 start_latent_idx = sum(self.num_latents_per_head[:head_idx])423 end_latent_idx = sum(self.num_latents_per_head[:head_idx + 1])424 split_size = self.latent_dim[start_latent_idx:end_latent_idx]425 outputs.extend(torch.split(codes, split_size, dim=1))426 max_dim = max(self.latent_dim)427 for i, dim in enumerate(self.latent_dim):428 if dim < max_dim:429 outputs[i] = F.pad(outputs[i], (0, max_dim - dim))430 outputs[i] = outputs[i].unsqueeze(1)431 432 return torch.cat(outputs, dim=1)433 434 435class FPN(nn.Module):436 """Implementation of Feature Pyramid Network (FPN).437 438 The input of this module is a pyramid of features with reducing resolutions.439 Then, this module fuses these multi-level features from `top_level` to440 `bottom_level`. In particular, starting from the `top_level`, each feature441 is convoluted, upsampled, and fused into its previous feature (which is also442 convoluted).443 444 Args:445 pyramid_channels: A list of integers, each of which indicates the number446 of channels of the feature from a particular level.447 out_channels: Number of channels for each output.448 449 Returns:450 A list of feature maps, each of which has `out_channels` channels.451 """452 453 def __init__(self, pyramid_channels, out_channels):454 super().__init__()455 assert isinstance(pyramid_channels, (list, tuple))456 self.num_levels = len(pyramid_channels)457 458 self.lateral_layers = nn.ModuleList()459 self.feature_layers = nn.ModuleList()460 for i in range(self.num_levels):461 in_channels = pyramid_channels[i]462 self.lateral_layers.append(nn.Conv2d(in_channels=in_channels,463 out_channels=out_channels,464 kernel_size=3,465 padding=1,466 bias=True))467 self.feature_layers.append(nn.Conv2d(in_channels=out_channels,468 out_channels=out_channels,469 kernel_size=3,470 padding=1,471 bias=True))472 473 def forward(self, inputs):474 if len(inputs) != self.num_levels:475 raise ValueError('Number of inputs and `num_levels` mismatch!')476 477 # Project all related features to `out_channels`.478 laterals = []479 for i in range(self.num_levels):480 laterals.append(self.lateral_layers[i](inputs[i]))481 482 # Fusion, starting from `top_level`.483 for i in range(self.num_levels - 1, 0, -1):484 scale_factor = laterals[i - 1].shape[2] // laterals[i].shape[2]485 laterals[i - 1] = (laterals[i - 1] +486 F.interpolate(laterals[i],487 mode='nearest',488 scale_factor=scale_factor))489 490 # Get outputs.491 outputs = []492 for i, lateral in enumerate(laterals):493 outputs.append(self.feature_layers[i](lateral))494 495 return outputs496 497 498class SAM(nn.Module):499 """Implementation of Spatial Alignment Module (SAM).500 501 The input of this module is a pyramid of features with reducing resolutions.502 Then this module downsamples all levels of feature to the minimum resolution503 and fuses it with the smallest feature map.504 505 Args:506 pyramid_channels: A list of integers, each of which indicates the number507 of channels of the feature from a particular level.508 out_channels: Number of channels for each output.509 510 Returns:511 A list of feature maps, each of which has `out_channels` channels.512 """513 514 def __init__(self, pyramid_channels, out_channels):515 super().__init__()516 assert isinstance(pyramid_channels, (list, tuple))517 self.num_levels = len(pyramid_channels)518 519 self.fusion_layers = nn.ModuleList()520 for i in range(self.num_levels):521 in_channels = pyramid_channels[i]522 self.fusion_layers.append(nn.Conv2d(in_channels=in_channels,523 out_channels=out_channels,524 kernel_size=3,525 padding=1,526 bias=True))527 528 def forward(self, inputs):529 if len(inputs) != self.num_levels:530 raise ValueError('Number of inputs and `num_levels` mismatch!')531 532 output_res = inputs[-1].shape[2:]533 for i in range(self.num_levels - 1, -1, -1):534 if i != self.num_levels - 1:535 inputs[i] = F.adaptive_avg_pool2d(inputs[i], output_res)536 inputs[i] = self.fusion_layers[i](inputs[i])537 if i != self.num_levels - 1:538 inputs[i] = inputs[i] + inputs[-1]539 540 return inputs541 542 543class CodeHead(nn.Module):544 """Implementation of the task-head to produce inverted codes."""545 546 def __init__(self, in_channels, out_channels, norm_layer):547 super().__init__()548 self.fc = nn.Linear(in_channels, out_channels, bias=True)549 if norm_layer is None:550 self.norm = nn.Identity()551 else:552 self.norm = norm_layer(out_channels)553 554 def forward(self, x):555 if x.ndim > 2:556 x = x.flatten(start_dim=1)557 latent = self.fc(x)558 latent = latent.unsqueeze(2).unsqueeze(3)559 latent = self.norm(latent)560 561 return latent.flatten(start_dim=1)562 563# pylint: enable=missing-function-docstring564