| |
| """ |
| Resnet based autoencoder models. |
| |
| File originally from https://github.com/Horizon2333/imagenet-autoencoder/blob/main/models/resnet.py. |
| |
| Modifications: |
| - Adding `sigmoid` argument so `nn.BCEWithLogitsLoss` can be used |
| - Z_channels argument to fingerprint size can be varied |
| - Create ResNetVAE class (which performed worse for clustering unfortunately). |
| """ |
|
|
| import torch |
| import torch.nn as nn |
|
|
| def BuildAutoEncoder(arch, sigmoid=False, z_channels=None): |
| if arch in ["resnet18", "resnet34", "resnet50", "resnet101", "resnet152"]: |
| configs, bottleneck = get_configs(arch) |
| return ResNetAutoEncoder(configs, bottleneck, sigmoid, z_channels=z_channels) |
| return None |
|
|
| def get_configs(arch='resnet50'): |
|
|
| |
|
|
| if arch == 'resnet18': |
| return [2, 2, 2, 2], False |
| elif arch == 'resnet34': |
| return [3, 4, 6, 3], False |
| elif arch == 'resnet50': |
| return [3, 4, 6, 3], True |
| elif arch == 'resnet101': |
| return [3, 4, 23, 3], True |
| elif arch == 'resnet152': |
| return [3, 8, 36, 3], True |
| else: |
| raise ValueError("Undefined model") |
|
|
| class ResNetAutoEncoder(nn.Module): |
|
|
| def __init__(self, configs, bottleneck, sigmoid, z_channels=None): |
|
|
| super(ResNetAutoEncoder, self).__init__() |
|
|
| self.encoder = ResNetEncoder(configs=configs, bottleneck=bottleneck, z_channels=z_channels) |
| self.decoder = ResNetDecoder(configs=configs[::-1], bottleneck=bottleneck, sigmoid=sigmoid, z_channels=z_channels) |
| |
| def forward(self, x): |
|
|
| x = self.encoder(x) |
| x = self.decoder(x) |
|
|
| return x |
| |
| class ResnetVAE(ResNetAutoEncoder): |
| def __init__(self, configs, bottleneck, sigmoid, z_channels): |
|
|
| super(ResnetVAE, self).__init__(configs, bottleneck, sigmoid) |
| |
| self.z_channels = z_channels |
| self.z_dim = z_channels * 4 * 4 |
| |
| |
| self.encoder = ResNetEncoder(configs=configs, bottleneck=bottleneck, z_channels=z_channels*2) |
| self.decoder = ResNetDecoder(configs=configs[::-1], bottleneck=bottleneck, sigmoid=sigmoid, z_channels=z_channels) |
| |
| self.flatten = nn.Flatten() |
| |
| def forward(self, x): |
| x = self.encoder(x) |
| mu_logvar = self.flatten(x) |
| mu = mu_logvar[:, :self.z_dim] |
| logvar = mu_logvar[:, self.z_dim:] |
| |
| z = self.reparametrize(mu, logvar) |
| res = z.view(z.shape[0], self.z_channels, 4, 4) |
| x_recon = self.decoder(res) |
| |
| return x_recon, mu, logvar |
|
|
| def reparametrize(self, mu, logvar): |
| std = torch.exp(0.5 * logvar) |
| eps = torch.randn_like(std) |
| return eps * std + mu |
| |
|
|
| class ResNet(nn.Module): |
| """ |
| Normal resnet for classification - not used |
| """ |
| def __init__(self, configs, bottleneck=False, num_classes=1000): |
| super(ResNet, self).__init__() |
|
|
| self.encoder = ResNetEncoder(configs, bottleneck) |
|
|
| self.avpool = nn.AdaptiveAvgPool2d((1,1)) |
|
|
| if bottleneck: |
| self.fc = nn.Linear(in_features=2048, out_features=num_classes) |
| else: |
| self.fc = nn.Linear(in_features=512, out_features=num_classes) |
|
|
| for m in self.modules(): |
| if isinstance(m, nn.Conv2d): |
| nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu") |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.BatchNorm2d): |
| nn.init.constant_(m.weight, 1) |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.Linear): |
| nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu") |
| nn.init.constant_(m.bias, 0) |
| |
| def forward(self, x): |
|
|
| x = self.encoder(x) |
|
|
| x = self.avpool(x) |
|
|
| x = torch.flatten(x, 1) |
|
|
| x = self.fc(x) |
|
|
| return x |
|
|
|
|
| class ResNetEncoder(nn.Module): |
|
|
| def __init__(self, configs, bottleneck=False, z_channels=None): |
| super(ResNetEncoder, self).__init__() |
|
|
| if len(configs) != 4: |
| raise ValueError("Only 4 layers can be configued") |
|
|
| self.conv1 = nn.Sequential( |
| nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=2, padding=3, bias=False), |
| nn.BatchNorm2d(num_features=64), |
| nn.ReLU(inplace=True), |
| ) |
| |
| if not z_channels: |
| if bottleneck: z_channels = 2048 |
| else: z_channels = 512 |
| |
| if bottleneck: |
|
|
| self.conv2 = EncoderBottleneckBlock(in_channels=64, hidden_channels=64, up_channels=256, layers=configs[0], downsample_method="pool") |
| self.conv3 = EncoderBottleneckBlock(in_channels=256, hidden_channels=128, up_channels=512, layers=configs[1], downsample_method="conv") |
| self.conv4 = EncoderBottleneckBlock(in_channels=512, hidden_channels=256, up_channels=1024, layers=configs[2], downsample_method="conv") |
| self.conv5 = EncoderBottleneckBlock(in_channels=1024, hidden_channels=512, up_channels=z_channels, layers=configs[3], downsample_method="conv") |
|
|
| else: |
|
|
| self.conv2 = EncoderResidualBlock(in_channels=64, hidden_channels=64, layers=configs[0], downsample_method="pool") |
| self.conv3 = EncoderResidualBlock(in_channels=64, hidden_channels=128, layers=configs[1], downsample_method="conv") |
| self.conv4 = EncoderResidualBlock(in_channels=128, hidden_channels=256, layers=configs[2], downsample_method="conv") |
| self.conv5 = EncoderResidualBlock(in_channels=256, hidden_channels=z_channels, layers=configs[3], downsample_method="conv") |
|
|
| def forward(self, x): |
|
|
| x = self.conv1(x) |
| x = self.conv2(x) |
| x = self.conv3(x) |
| x = self.conv4(x) |
| x = self.conv5(x) |
|
|
| return x |
| |
| class ResNetDecoder(nn.Module): |
|
|
| def __init__(self, configs, bottleneck=False, sigmoid=False, z_channels=None): |
| super(ResNetDecoder, self).__init__() |
|
|
| if len(configs) != 4: |
| raise ValueError("Only 4 layers can be configued") |
| |
| if not z_channels: |
| if bottleneck: z_channels = 2048 |
| else: z_channels = 512 |
|
|
| if bottleneck: |
|
|
| self.conv1 = DecoderBottleneckBlock(in_channels=z_channels, hidden_channels=512, down_channels=1024, layers=configs[0]) |
| self.conv2 = DecoderBottleneckBlock(in_channels=1024, hidden_channels=256, down_channels=512, layers=configs[1]) |
| self.conv3 = DecoderBottleneckBlock(in_channels=512, hidden_channels=128, down_channels=256, layers=configs[2]) |
| self.conv4 = DecoderBottleneckBlock(in_channels=256, hidden_channels=64, down_channels=64, layers=configs[3]) |
|
|
|
|
| else: |
|
|
| self.conv1 = DecoderResidualBlock(hidden_channels=z_channels, output_channels=256, layers=configs[0]) |
| self.conv2 = DecoderResidualBlock(hidden_channels=256, output_channels=128, layers=configs[1]) |
| self.conv3 = DecoderResidualBlock(hidden_channels=128, output_channels=64, layers=configs[2]) |
| self.conv4 = DecoderResidualBlock(hidden_channels=64, output_channels=64, layers=configs[3]) |
|
|
| self.conv5 = nn.Sequential( |
| nn.BatchNorm2d(num_features=64), |
| nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=7, stride=2, padding=3, output_padding=1, bias=False), |
| ) |
| |
| if sigmoid: |
| self.gate = nn.Sigmoid() |
| else: |
| self.gate = nn.Identity() |
|
|
| def forward(self, x): |
| |
| x = self.conv1(x) |
| x = self.conv2(x) |
| x = self.conv3(x) |
| x = self.conv4(x) |
| x = self.conv5(x) |
| x = self.gate(x) |
|
|
| return x |
| |
| class EncoderResidualBlock(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, layers, downsample_method="conv"): |
| super(EncoderResidualBlock, self).__init__() |
|
|
| if downsample_method == "conv": |
|
|
| for i in range(layers): |
|
|
| if i == 0: |
| layer = EncoderResidualLayer(in_channels=in_channels, hidden_channels=hidden_channels, downsample=True) |
| else: |
| layer = EncoderResidualLayer(in_channels=hidden_channels, hidden_channels=hidden_channels, downsample=False) |
| |
| self.add_module('%02d EncoderLayer' % i, layer) |
| |
| elif downsample_method == "pool": |
|
|
| maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
|
|
| self.add_module('00 MaxPooling', maxpool) |
|
|
| for i in range(layers): |
|
|
| if i == 0: |
| layer = EncoderResidualLayer(in_channels=in_channels, hidden_channels=hidden_channels, downsample=False) |
| else: |
| layer = EncoderResidualLayer(in_channels=hidden_channels, hidden_channels=hidden_channels, downsample=False) |
| |
| self.add_module('%02d EncoderLayer' % (i+1), layer) |
| |
| def forward(self, x): |
|
|
| for name, layer in self.named_children(): |
|
|
| x = layer(x) |
|
|
| return x |
|
|
| class EncoderBottleneckBlock(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, up_channels, layers, downsample_method="conv"): |
| super(EncoderBottleneckBlock, self).__init__() |
|
|
| if downsample_method == "conv": |
|
|
| for i in range(layers): |
|
|
| if i == 0: |
| layer = EncoderBottleneckLayer(in_channels=in_channels, hidden_channels=hidden_channels, up_channels=up_channels, downsample=True) |
| else: |
| layer = EncoderBottleneckLayer(in_channels=up_channels, hidden_channels=hidden_channels, up_channels=up_channels, downsample=False) |
| |
| self.add_module('%02d EncoderLayer' % i, layer) |
| |
| elif downsample_method == "pool": |
|
|
| maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
|
|
| self.add_module('00 MaxPooling', maxpool) |
|
|
| for i in range(layers): |
|
|
| if i == 0: |
| layer = EncoderBottleneckLayer(in_channels=in_channels, hidden_channels=hidden_channels, up_channels=up_channels, downsample=False) |
| else: |
| layer = EncoderBottleneckLayer(in_channels=up_channels, hidden_channels=hidden_channels, up_channels=up_channels, downsample=False) |
| |
| self.add_module('%02d EncoderLayer' % (i+1), layer) |
| |
| def forward(self, x): |
|
|
| for name, layer in self.named_children(): |
|
|
| x = layer(x) |
|
|
| return x |
|
|
| class DecoderResidualBlock(nn.Module): |
|
|
| def __init__(self, hidden_channels, output_channels, layers): |
| super(DecoderResidualBlock, self).__init__() |
|
|
| for i in range(layers): |
|
|
| if i == layers - 1: |
| layer = DecoderResidualLayer(hidden_channels=hidden_channels, output_channels=output_channels, upsample=True) |
| else: |
| layer = DecoderResidualLayer(hidden_channels=hidden_channels, output_channels=hidden_channels, upsample=False) |
| |
| self.add_module('%02d EncoderLayer' % i, layer) |
| |
| def forward(self, x): |
|
|
| for name, layer in self.named_children(): |
|
|
| x = layer(x) |
|
|
| return x |
|
|
| class DecoderBottleneckBlock(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, down_channels, layers): |
| super(DecoderBottleneckBlock, self).__init__() |
|
|
| for i in range(layers): |
|
|
| if i == layers - 1: |
| layer = DecoderBottleneckLayer(in_channels=in_channels, hidden_channels=hidden_channels, down_channels=down_channels, upsample=True) |
| else: |
| layer = DecoderBottleneckLayer(in_channels=in_channels, hidden_channels=hidden_channels, down_channels=in_channels, upsample=False) |
| |
| self.add_module('%02d EncoderLayer' % i, layer) |
| |
| |
| def forward(self, x): |
|
|
| for name, layer in self.named_children(): |
|
|
| x = layer(x) |
|
|
| return x |
|
|
| class EncoderResidualLayer(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, downsample): |
| super(EncoderResidualLayer, self).__init__() |
|
|
| if downsample: |
| self.weight_layer1 = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=3, stride=2, padding=1, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| ) |
| else: |
| self.weight_layer1 = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| ) |
|
|
| self.weight_layer2 = nn.Sequential( |
| nn.Conv2d(in_channels=hidden_channels, out_channels=hidden_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| ) |
|
|
| if downsample: |
| self.downsample = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=1, stride=2, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| ) |
| else: |
| self.downsample = None |
|
|
| self.relu = nn.Sequential( |
| nn.ReLU(inplace=True) |
| ) |
| |
| def forward(self, x): |
|
|
| identity = x |
|
|
| x = self.weight_layer1(x) |
| x = self.weight_layer2(x) |
|
|
| if self.downsample is not None: |
| identity = self.downsample(identity) |
|
|
| x = x + identity |
|
|
| x = self.relu(x) |
|
|
| return x |
|
|
| class EncoderBottleneckLayer(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, up_channels, downsample): |
| super(EncoderBottleneckLayer, self).__init__() |
|
|
| if downsample: |
| self.weight_layer1 = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=1, stride=2, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| ) |
| else: |
| self.weight_layer1 = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=1, stride=1, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| ) |
|
|
| self.weight_layer2 = nn.Sequential( |
| nn.Conv2d(in_channels=hidden_channels, out_channels=hidden_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| ) |
|
|
| self.weight_layer3 = nn.Sequential( |
| nn.Conv2d(in_channels=hidden_channels, out_channels=up_channels, kernel_size=1, stride=1, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=up_channels), |
| ) |
|
|
| if downsample: |
| self.downsample = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=up_channels, kernel_size=1, stride=2, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=up_channels), |
| ) |
| elif (in_channels != up_channels): |
| self.downsample = None |
| self.up_scale = nn.Sequential( |
| nn.Conv2d(in_channels=in_channels, out_channels=up_channels, kernel_size=1, stride=1, padding=0, bias=False), |
| nn.BatchNorm2d(num_features=up_channels), |
| ) |
| else: |
| self.downsample = None |
| self.up_scale = None |
|
|
| self.relu = nn.Sequential( |
| nn.ReLU(inplace=True) |
| ) |
| |
| def forward(self, x): |
|
|
| identity = x |
|
|
| x = self.weight_layer1(x) |
| x = self.weight_layer2(x) |
| x = self.weight_layer3(x) |
|
|
| if self.downsample is not None: |
| identity = self.downsample(identity) |
| elif self.up_scale is not None: |
| identity = self.up_scale(identity) |
|
|
| x = x + identity |
|
|
| x = self.relu(x) |
|
|
| return x |
|
|
| class DecoderResidualLayer(nn.Module): |
|
|
| def __init__(self, hidden_channels, output_channels, upsample): |
| super(DecoderResidualLayer, self).__init__() |
|
|
| self.weight_layer1 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=hidden_channels, out_channels=hidden_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| ) |
|
|
| if upsample: |
| self.weight_layer2 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(in_channels=hidden_channels, out_channels=output_channels, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False) |
| ) |
| else: |
| self.weight_layer2 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=hidden_channels, out_channels=output_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| ) |
|
|
| if upsample: |
| self.upsample = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(in_channels=hidden_channels, out_channels=output_channels, kernel_size=1, stride=2, output_padding=1, bias=False) |
| ) |
| else: |
| self.upsample = None |
| |
| def forward(self, x): |
|
|
| identity = x |
|
|
| x = self.weight_layer1(x) |
| x = self.weight_layer2(x) |
|
|
| if self.upsample is not None: |
| identity = self.upsample(identity) |
|
|
| x = x + identity |
|
|
| return x |
|
|
| class DecoderBottleneckLayer(nn.Module): |
|
|
| def __init__(self, in_channels, hidden_channels, down_channels, upsample): |
| super(DecoderBottleneckLayer, self).__init__() |
|
|
| self.weight_layer1 = nn.Sequential( |
| nn.BatchNorm2d(num_features=in_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=1, stride=1, padding=0, bias=False), |
| ) |
|
|
| self.weight_layer2 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=hidden_channels, out_channels=hidden_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| ) |
|
|
| if upsample: |
| self.weight_layer3 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(in_channels=hidden_channels, out_channels=down_channels, kernel_size=1, stride=2, output_padding=1, bias=False) |
| ) |
| else: |
| self.weight_layer3 = nn.Sequential( |
| nn.BatchNorm2d(num_features=hidden_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=hidden_channels, out_channels=down_channels, kernel_size=1, stride=1, padding=0, bias=False) |
| ) |
|
|
| if upsample: |
| self.upsample = nn.Sequential( |
| nn.BatchNorm2d(num_features=in_channels), |
| nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(in_channels=in_channels, out_channels=down_channels, kernel_size=1, stride=2, output_padding=1, bias=False) |
| ) |
| elif (in_channels != down_channels): |
| self.upsample = None |
| self.down_scale = nn.Sequential( |
| nn.BatchNorm2d(num_features=in_channels), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(in_channels=in_channels, out_channels=down_channels, kernel_size=1, stride=1, padding=0, bias=False) |
| ) |
| else: |
| self.upsample = None |
| self.down_scale = None |
| |
| def forward(self, x): |
|
|
| identity = x |
|
|
| x = self.weight_layer1(x) |
| x = self.weight_layer2(x) |
| x = self.weight_layer3(x) |
|
|
| if self.upsample is not None: |
| identity = self.upsample(identity) |
| elif self.down_scale is not None: |
| identity = self.down_scale(identity) |
|
|
| x = x + identity |
|
|
| return x |
| |
| |
| class ResidualLayer(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.conv = nn.Conv2d(32, 32, 1) |
| def forward(self, x): |
| return x + self.conv(x) |
|
|
| if __name__ == "__main__": |
|
|
| configs, bottleneck = get_configs("resnet152") |
|
|
| encoder = ResNetEncoder(configs, bottleneck) |
|
|
| input = torch.randn((5,3,224,224)) |
|
|
| print(input.shape) |
|
|
| output = encoder(input) |
|
|
| print(output.shape) |
|
|
| decoder = ResNetDecoder(configs[::-1], bottleneck) |
|
|
| output = decoder(output) |
|
|
| print(output.shape) |
