fix error

tools/visualization_0416/utils/model_0506/model/basic_model/basic_block.py

@@ -0,0 +1,163 @@
+"""
+Reference: https://github.com/hedra-labs/one-shot-face/blob/mega-portraits/models/building_blocks.py
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import spectral_norm
+
+USE_BIAS = False
+
+# https://github.com/joe-siyuan-qiao/WeightStandardization?tab=readme-ov-file#pytorch
+class WSConv2d(nn.Conv2d):
+    def __init__(self, *args, **kwargs):
+        super(WSConv2d, self).__init__(*args, **kwargs)
+
+    def forward(self, inp):
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True)
+        weight = weight - weight_mean
+        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv2d(inp, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class WSConv3d(nn.Conv3d):
+    def __init__(self, *args, **kwargs):
+        super(WSConv3d, self).__init__(*args, **kwargs)
+
+    def forward(self, inp):
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True).mean(dim=4, keepdim=True)
+        weight = weight - weight_mean
+        std = weight.view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv3d(inp, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class ResBlock2d(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, num_channels_per_group: int, use_spectral_norm: bool = False):
+        super().__init__()
+
+        norm_func = lambda x: x
+        if use_spectral_norm:
+            norm_func = spectral_norm
+
+        if in_channels != out_channels:
+            self.skip_layer = norm_func(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS))
+        else:
+            self.skip_layer = lambda x: x
+
+        self.layers = nn.Sequential(
+            nn.GroupNorm(in_channels // num_channels_per_group, in_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            norm_func(WSConv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS)),
+            nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            norm_func(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS)),
+        )
+
+    def forward(self, inp: torch.Tensor):
+        return self.skip_layer(inp) + self.layers(inp)
+
+
+class ResBlock3d(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, num_channels_per_group: int):
+        super().__init__()
+
+        if in_channels != out_channels:
+            self.skip_layer = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS)
+        else:
+            self.skip_layer = lambda x: x
+
+        self.layers = nn.Sequential(
+            nn.GroupNorm(in_channels // num_channels_per_group, in_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            WSConv3d(in_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS),
+            nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1, bias=USE_BIAS),
+        )
+
+    def forward(self, inp: torch.Tensor):
+        return self.skip_layer(inp) + self.layers(inp)
+
+
+class ResBasic(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, stride: int, num_channels_per_group: int):
+        super().__init__()
+
+        if stride != 1 and stride != 2:
+            raise NotImplementedError(f"Stride can be only 1 or 2 but '{stride}' is passed.")
+
+        if in_channels != out_channels or stride != 1:
+            self.skip_layer = nn.Sequential(
+                WSConv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=USE_BIAS),
+                nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+            )
+        else:
+            self.skip_layer = lambda x: x
+
+        self.layers = nn.Sequential(
+            WSConv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=USE_BIAS),
+            nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            WSConv2d(out_channels, out_channels, kernel_size=1, bias=USE_BIAS),
+            nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+        )
+
+
+    def forward(self, inp: torch.Tensor):
+        return F.relu(self.skip_layer(inp) + self.layers(inp))
+
+
+class ResBottleneck(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, stride: int, num_channels_per_group: int):
+        super().__init__()
+
+        if stride != 1 and stride != 2:
+            raise NotImplementedError(f"Stride can be only 1 or 2 but '{stride}' is passed.")
+
+        if in_channels != out_channels or stride != 1:
+            self.skip_layer = nn.Sequential(
+                WSConv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=USE_BIAS),
+                nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+            )
+        else:
+            self.skip_layer = lambda x: x
+
+        temp_out_channels = out_channels // 4
+        self.layers = nn.Sequential(
+            WSConv2d(in_channels, temp_out_channels, kernel_size=1, bias=USE_BIAS),
+            nn.GroupNorm(temp_out_channels // num_channels_per_group, temp_out_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            WSConv2d(temp_out_channels, temp_out_channels, kernel_size=3, stride=stride, padding=1, bias=USE_BIAS),
+            nn.GroupNorm(temp_out_channels // num_channels_per_group, temp_out_channels, affine=not USE_BIAS),
+            nn.ReLU(inplace=True),
+            WSConv2d(temp_out_channels, out_channels, kernel_size=1, bias=USE_BIAS),
+            nn.GroupNorm(out_channels // num_channels_per_group, out_channels, affine=not USE_BIAS),
+        )
+
+
+    def forward(self, inp: torch.Tensor):
+        return F.relu(self.skip_layer(inp) + self.layers(inp))
+
+
+class ReshapeTo3DLayer(nn.Module):
+    def __init__(self, out_depth: int):
+        super().__init__()
+
+        self.out_depth = out_depth
+
+    def forward(self, inp: torch.Tensor):
+        batch_size, channels, height, width = inp.shape
+        return inp.view(batch_size, channels // self.out_depth, self.out_depth, height, width)
+
+
+class ReshapeTo2DLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, inp: torch.Tensor):
+        batch_size, channels, depth, height, width = inp.shape
+        return inp.view(batch_size, channels * depth, height, width)

tools/visualization_0416/utils/model_0506/model/basic_model/frnet.py

@@ -0,0 +1,177 @@
+"""
+Reference: https://github.com/yfeng95/DECA/blob/a11554ae2a2b0f3998cf1fa94dd4db03babb34a2/decalib/models/frnet.py
+"""
+import torch.nn as nn
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import math
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, include_top=True):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.include_top = include_top
+        
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        
+        if not self.include_top:
+            return x
+        
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+def resnet50(**kwargs):
+    """Constructs a ResNet-50 model.
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    return model
+
+import pickle
+def load_state_dict(model, fname):
+    """
+    Set parameters converted from Caffe models authors of VGGFace2 provide.
+    See https://www.robots.ox.ac.uk/~vgg/data/vgg_face2/.
+    Arguments:
+        model: model
+        fname: file name of parameters converted from a Caffe model, assuming the file format is Pickle.
+    """
+    with open(fname, 'rb') as f:
+        weights = pickle.load(f, encoding='latin1')
+
+    own_state = model.state_dict()
+    for name, param in weights.items():
+        if name in own_state:
+            try:
+                own_state[name].copy_(torch.from_numpy(param))
+            except Exception:
+                raise RuntimeError('While copying the parameter named {}, whose dimensions in the model are {} and whose '\
+                                   'dimensions in the checkpoint are {}.'.format(name, own_state[name].size(), param.size()))
+        else:
+            raise KeyError('unexpected key "{}" in state_dict'.format(name))

tools/visualization_0416/utils/model_0506/model/basic_model/resnet.py

@@ -0,0 +1,182 @@
+'''ResNet in PyTorch.
+
+For Pre-activation ResNet, see 'preact_resnet.py'.
+
+Reference:
+[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
+    Deep Residual Learning for Image Recognition. arXiv:1512.03385
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion *
+                               planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=7):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        # self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv1 = nn.Conv2d(1, 64, kernel_size = 3, stride = 1, padding = 1, bias = False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU()
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512*block.expansion, num_classes)
+        self.fc = nn.Linear(512*block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+class ResNet_AE(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=7):
+        super(ResNet_AE, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512*block.expansion, num_classes)
+
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(64, 1, kernel_size=3,stride=1, padding=1, bias=False),
+            nn.Sigmoid()
+        )
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+        out1 = self.conv1(x)
+        
+        decoded = self.decoder(out1)
+
+        out = F.relu(self.bn1(out1))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out, x, decoded
+
+
+def ResNet18_AE():
+    return ResNet_AE(BasicBlock, [2, 2, 2, 2])
+
+
+def ResNet18():
+    return ResNet(BasicBlock, [2, 2, 2, 2])
+
+
+def ResNet34():
+    return ResNet(BasicBlock, [3, 4, 6, 3])
+
+
+def ResNet50():
+    return ResNet(Bottleneck, [3, 4, 6, 3])
+
+
+def ResNet101():
+    return ResNet(Bottleneck, [3, 4, 23, 3])
+
+
+def ResNet152():
+    return ResNet(Bottleneck, [3, 8, 36, 3])
+
+
+def test():
+    net = ResNet18()
+    y = net(torch.randn(1, 3, 32, 32))
+    print(y.size())
+
+# test()