RabbitRUI/ruispace
0
1"""This script defines deep neural networks for Deep3DFaceRecon_pytorch2"""3 4import os5import numpy as np6import torch.nn.functional as F7from torch.nn import init8import functools9from torch.optim import lr_scheduler10import torch11from torch import Tensor12import torch.nn as nn13try:14 from torch.hub import load_state_dict_from_url15except ImportError:16 from torch.utils.model_zoo import load_url as load_state_dict_from_url17from typing import Type, Any, Callable, Union, List, Optional18from .arcface_torch.backbones import get_model19from kornia.geometry import warp_affine20 21def resize_n_crop(image, M, dsize=112):22 # image: (b, c, h, w)23 # M : (b, 2, 3)24 return warp_affine(image, M, dsize=(dsize, dsize), align_corners=True)25 26def filter_state_dict(state_dict, remove_name='fc'):27 new_state_dict = {}28 for key in state_dict:29 if remove_name in key:30 continue31 new_state_dict[key] = state_dict[key]32 return new_state_dict33 34def get_scheduler(optimizer, opt):35 """Return a learning rate scheduler36 37 Parameters:38 optimizer -- the optimizer of the network39 opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions． 40 opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine41 42 For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.43 See https://pytorch.org/docs/stable/optim.html for more details.44 """45 if opt.lr_policy == 'linear':46 def lambda_rule(epoch):47 lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs + 1)48 return lr_l49 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)50 elif opt.lr_policy == 'step':51 scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_epochs, gamma=0.2)52 elif opt.lr_policy == 'plateau':53 scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)54 elif opt.lr_policy == 'cosine':55 scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)56 else:57 return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)58 return scheduler59 60 61def define_net_recon(net_recon, use_last_fc=False, init_path=None):62 return ReconNetWrapper(net_recon, use_last_fc=use_last_fc, init_path=init_path)63 64def define_net_recog(net_recog, pretrained_path=None):65 net = RecogNetWrapper(net_recog=net_recog, pretrained_path=pretrained_path)66 net.eval()67 return net68 69class ReconNetWrapper(nn.Module):70 fc_dim=25771 def __init__(self, net_recon, use_last_fc=False, init_path=None):72 super(ReconNetWrapper, self).__init__()73 self.use_last_fc = use_last_fc74 if net_recon not in func_dict:75 return NotImplementedError('network [%s] is not implemented', net_recon)76 func, last_dim = func_dict[net_recon]77 backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_dim)78 if init_path and os.path.isfile(init_path):79 state_dict = filter_state_dict(torch.load(init_path, map_location='cpu'))80 backbone.load_state_dict(state_dict)81 print("loading init net_recon %s from %s" %(net_recon, init_path))82 self.backbone = backbone83 if not use_last_fc:84 self.final_layers = nn.ModuleList([85 conv1x1(last_dim, 80, bias=True), # id layer86 conv1x1(last_dim, 64, bias=True), # exp layer87 conv1x1(last_dim, 80, bias=True), # tex layer88 conv1x1(last_dim, 3, bias=True), # angle layer89 conv1x1(last_dim, 27, bias=True), # gamma layer90 conv1x1(last_dim, 2, bias=True), # tx, ty91 conv1x1(last_dim, 1, bias=True) # tz92 ])93 for m in self.final_layers:94 nn.init.constant_(m.weight, 0.)95 nn.init.constant_(m.bias, 0.)96 97 def forward(self, x):98 x = self.backbone(x)99 if not self.use_last_fc:100 output = []101 for layer in self.final_layers:102 output.append(layer(x))103 x = torch.flatten(torch.cat(output, dim=1), 1)104 return x105 106 107class RecogNetWrapper(nn.Module):108 def __init__(self, net_recog, pretrained_path=None, input_size=112):109 super(RecogNetWrapper, self).__init__()110 net = get_model(name=net_recog, fp16=False)111 if pretrained_path:112 state_dict = torch.load(pretrained_path, map_location='cpu')113 net.load_state_dict(state_dict)114 print("loading pretrained net_recog %s from %s" %(net_recog, pretrained_path))115 for param in net.parameters():116 param.requires_grad = False117 self.net = net118 self.preprocess = lambda x: 2 * x - 1119 self.input_size=input_size120 121 def forward(self, image, M):122 image = self.preprocess(resize_n_crop(image, M, self.input_size))123 id_feature = F.normalize(self.net(image), dim=-1, p=2)124 return id_feature125 126 127# adapted from https://github.com/pytorch/vision/edit/master/torchvision/models/resnet.py128__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',129 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',130 'wide_resnet50_2', 'wide_resnet101_2']131 132 133model_urls = {134 'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth',135 'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth',136 'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',137 'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth',138 'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth',139 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',140 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',141 'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',142 'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',143}144 145 146def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:147 """3x3 convolution with padding"""148 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,149 padding=dilation, groups=groups, bias=False, dilation=dilation)150 151 152def conv1x1(in_planes: int, out_planes: int, stride: int = 1, bias: bool = False) -> nn.Conv2d:153 """1x1 convolution"""154 return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=bias)155 156 157class BasicBlock(nn.Module):158 expansion: int = 1159 160 def __init__(161 self,162 inplanes: int,163 planes: int,164 stride: int = 1,165 downsample: Optional[nn.Module] = None,166 groups: int = 1,167 base_width: int = 64,168 dilation: int = 1,169 norm_layer: Optional[Callable[..., nn.Module]] = None170 ) -> None:171 super(BasicBlock, self).__init__()172 if norm_layer is None:173 norm_layer = nn.BatchNorm2d174 if groups != 1 or base_width != 64:175 raise ValueError('BasicBlock only supports groups=1 and base_width=64')176 if dilation > 1:177 raise NotImplementedError("Dilation > 1 not supported in BasicBlock")178 # Both self.conv1 and self.downsample layers downsample the input when stride != 1179 self.conv1 = conv3x3(inplanes, planes, stride)180 self.bn1 = norm_layer(planes)181 self.relu = nn.ReLU(inplace=True)182 self.conv2 = conv3x3(planes, planes)183 self.bn2 = norm_layer(planes)184 self.downsample = downsample185 self.stride = stride186 187 def forward(self, x: Tensor) -> Tensor:188 identity = x189 190 out = self.conv1(x)191 out = self.bn1(out)192 out = self.relu(out)193 194 out = self.conv2(out)195 out = self.bn2(out)196 197 if self.downsample is not None:198 identity = self.downsample(x)199 200 out += identity201 out = self.relu(out)202 203 return out204 205 206class Bottleneck(nn.Module):207 # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)208 # while original implementation places the stride at the first 1x1 convolution(self.conv1)209 # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.210 # This variant is also known as ResNet V1.5 and improves accuracy according to211 # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.212 213 expansion: int = 4214 215 def __init__(216 self,217 inplanes: int,218 planes: int,219 stride: int = 1,220 downsample: Optional[nn.Module] = None,221 groups: int = 1,222 base_width: int = 64,223 dilation: int = 1,224 norm_layer: Optional[Callable[..., nn.Module]] = None225 ) -> None:226 super(Bottleneck, self).__init__()227 if norm_layer is None:228 norm_layer = nn.BatchNorm2d229 width = int(planes * (base_width / 64.)) * groups230 # Both self.conv2 and self.downsample layers downsample the input when stride != 1231 self.conv1 = conv1x1(inplanes, width)232 self.bn1 = norm_layer(width)233 self.conv2 = conv3x3(width, width, stride, groups, dilation)234 self.bn2 = norm_layer(width)235 self.conv3 = conv1x1(width, planes * self.expansion)236 self.bn3 = norm_layer(planes * self.expansion)237 self.relu = nn.ReLU(inplace=True)238 self.downsample = downsample239 self.stride = stride240 241 def forward(self, x: Tensor) -> Tensor:242 identity = x243 244 out = self.conv1(x)245 out = self.bn1(out)246 out = self.relu(out)247 248 out = self.conv2(out)249 out = self.bn2(out)250 out = self.relu(out)251 252 out = self.conv3(out)253 out = self.bn3(out)254 255 if self.downsample is not None:256 identity = self.downsample(x)257 258 out += identity259 out = self.relu(out)260 261 return out262 263 264class ResNet(nn.Module):265 266 def __init__(267 self,268 block: Type[Union[BasicBlock, Bottleneck]],269 layers: List[int],270 num_classes: int = 1000,271 zero_init_residual: bool = False,272 use_last_fc: bool = False,273 groups: int = 1,274 width_per_group: int = 64,275 replace_stride_with_dilation: Optional[List[bool]] = None,276 norm_layer: Optional[Callable[..., nn.Module]] = None277 ) -> None:278 super(ResNet, self).__init__()279 if norm_layer is None:280 norm_layer = nn.BatchNorm2d281 self._norm_layer = norm_layer282 283 self.inplanes = 64284 self.dilation = 1285 if replace_stride_with_dilation is None:286 # each element in the tuple indicates if we should replace287 # the 2x2 stride with a dilated convolution instead288 replace_stride_with_dilation = [False, False, False]289 if len(replace_stride_with_dilation) != 3:290 raise ValueError("replace_stride_with_dilation should be None "291 "or a 3-element tuple, got {}".format(replace_stride_with_dilation))292 self.use_last_fc = use_last_fc293 self.groups = groups294 self.base_width = width_per_group295 self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,296 bias=False)297 self.bn1 = norm_layer(self.inplanes)298 self.relu = nn.ReLU(inplace=True)299 self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)300 self.layer1 = self._make_layer(block, 64, layers[0])301 self.layer2 = self._make_layer(block, 128, layers[1], stride=2,302 dilate=replace_stride_with_dilation[0])303 self.layer3 = self._make_layer(block, 256, layers[2], stride=2,304 dilate=replace_stride_with_dilation[1])305 self.layer4 = self._make_layer(block, 512, layers[3], stride=2,306 dilate=replace_stride_with_dilation[2])307 self.avgpool = nn.AdaptiveAvgPool2d((1, 1))308 309 if self.use_last_fc:310 self.fc = nn.Linear(512 * block.expansion, num_classes)311 312 for m in self.modules():313 if isinstance(m, nn.Conv2d):314 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')315 elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):316 nn.init.constant_(m.weight, 1)317 nn.init.constant_(m.bias, 0)318 319 320 321 # Zero-initialize the last BN in each residual branch,322 # so that the residual branch starts with zeros, and each residual block behaves like an identity.323 # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677324 if zero_init_residual:325 for m in self.modules():326 if isinstance(m, Bottleneck):327 nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type]328 elif isinstance(m, BasicBlock):329 nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type]330 331 def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,332 stride: int = 1, dilate: bool = False) -> nn.Sequential:333 norm_layer = self._norm_layer334 downsample = None335 previous_dilation = self.dilation336 if dilate:337 self.dilation *= stride338 stride = 1339 if stride != 1 or self.inplanes != planes * block.expansion:340 downsample = nn.Sequential(341 conv1x1(self.inplanes, planes * block.expansion, stride),342 norm_layer(planes * block.expansion),343 )344 345 layers = []346 layers.append(block(self.inplanes, planes, stride, downsample, self.groups,347 self.base_width, previous_dilation, norm_layer))348 self.inplanes = planes * block.expansion349 for _ in range(1, blocks):350 layers.append(block(self.inplanes, planes, groups=self.groups,351 base_width=self.base_width, dilation=self.dilation,352 norm_layer=norm_layer))353 354 return nn.Sequential(*layers)355 356 def _forward_impl(self, x: Tensor) -> Tensor:357 # See note [TorchScript super()]358 x = self.conv1(x)359 x = self.bn1(x)360 x = self.relu(x)361 x = self.maxpool(x)362 363 x = self.layer1(x)364 x = self.layer2(x)365 x = self.layer3(x)366 x = self.layer4(x)367 368 x = self.avgpool(x)369 if self.use_last_fc:370 x = torch.flatten(x, 1)371 x = self.fc(x)372 return x373 374 def forward(self, x: Tensor) -> Tensor:375 return self._forward_impl(x)376 377 378def _resnet(379 arch: str,380 block: Type[Union[BasicBlock, Bottleneck]],381 layers: List[int],382 pretrained: bool,383 progress: bool,384 **kwargs: Any385) -> ResNet:386 model = ResNet(block, layers, **kwargs)387 if pretrained:388 state_dict = load_state_dict_from_url(model_urls[arch],389 progress=progress)390 model.load_state_dict(state_dict)391 return model392 393 394def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:395 r"""ResNet-18 model from396 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.397 398 Args:399 pretrained (bool): If True, returns a model pre-trained on ImageNet400 progress (bool): If True, displays a progress bar of the download to stderr401 """402 return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,403 **kwargs)404 405 406def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:407 r"""ResNet-34 model from408 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.409 410 Args:411 pretrained (bool): If True, returns a model pre-trained on ImageNet412 progress (bool): If True, displays a progress bar of the download to stderr413 """414 return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,415 **kwargs)416 417 418def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:419 r"""ResNet-50 model from420 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.421 422 Args:423 pretrained (bool): If True, returns a model pre-trained on ImageNet424 progress (bool): If True, displays a progress bar of the download to stderr425 """426 return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,427 **kwargs)428 429 430def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:431 r"""ResNet-101 model from432 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.433 434 Args:435 pretrained (bool): If True, returns a model pre-trained on ImageNet436 progress (bool): If True, displays a progress bar of the download to stderr437 """438 return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,439 **kwargs)440 441 442def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:443 r"""ResNet-152 model from444 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.445 446 Args:447 pretrained (bool): If True, returns a model pre-trained on ImageNet448 progress (bool): If True, displays a progress bar of the download to stderr449 """450 return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,451 **kwargs)452 453 454def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:455 r"""ResNeXt-50 32x4d model from456 `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.457 458 Args:459 pretrained (bool): If True, returns a model pre-trained on ImageNet460 progress (bool): If True, displays a progress bar of the download to stderr461 """462 kwargs['groups'] = 32463 kwargs['width_per_group'] = 4464 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],465 pretrained, progress, **kwargs)466 467 468def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:469 r"""ResNeXt-101 32x8d model from470 `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.471 472 Args:473 pretrained (bool): If True, returns a model pre-trained on ImageNet474 progress (bool): If True, displays a progress bar of the download to stderr475 """476 kwargs['groups'] = 32477 kwargs['width_per_group'] = 8478 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],479 pretrained, progress, **kwargs)480 481 482def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:483 r"""Wide ResNet-50-2 model from484 `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.485 486 The model is the same as ResNet except for the bottleneck number of channels487 which is twice larger in every block. The number of channels in outer 1x1488 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048489 channels, and in Wide ResNet-50-2 has 2048-1024-2048.490 491 Args:492 pretrained (bool): If True, returns a model pre-trained on ImageNet493 progress (bool): If True, displays a progress bar of the download to stderr494 """495 kwargs['width_per_group'] = 64 * 2496 return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],497 pretrained, progress, **kwargs)498 499 500def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:501 r"""Wide ResNet-101-2 model from502 `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.503 504 The model is the same as ResNet except for the bottleneck number of channels505 which is twice larger in every block. The number of channels in outer 1x1506 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048507 channels, and in Wide ResNet-50-2 has 2048-1024-2048.508 509 Args:510 pretrained (bool): If True, returns a model pre-trained on ImageNet511 progress (bool): If True, displays a progress bar of the download to stderr512 """513 kwargs['width_per_group'] = 64 * 2514 return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],515 pretrained, progress, **kwargs)516 517 518func_dict = {519 'resnet18': (resnet18, 512),520 'resnet50': (resnet50, 2048)521}522