CalvinBackup/SimSwapPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

eca depth wise

Browse Source
This commit is contained in:
chenxuanhong
2022-02-19 18:26:22 +08:00
parent e86270032b
commit db049166a0
7 changed files with 264 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/DeConv.py
+41
View File
@@ -0,0 +1,41 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: DeConv.py
# Created Date: Tuesday July 20th 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Saturday, 19th February 2022 5:35:38 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
from torch import nn
class DeConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size = 3, upsampl_scale = 2, padding="zero"):
super().__init__()
self.upsampling = nn.UpsamplingBilinear2d(scale_factor=upsampl_scale)
padding_size = int((kernel_size -1)/2)
if padding.lower() == "reflect":
self.conv = nn.Sequential(
nn.ReflectionPad2d(padding_size),
nn.Conv2d(in_channels = in_channels,
out_channels = out_channels, kernel_size= kernel_size, bias= False))
# for layer in self.conv:
# if isinstance(layer,nn.Conv2d):
# nn.init.xavier_uniform_(layer.weight)
elif padding.lower() == "zero":
self.conv = nn.Conv2d(in_channels = in_channels, padding = 1,
out_channels = out_channels, kernel_size= kernel_size, bias= False)
# nn.init.xavier_uniform_(self.conv.weight)
# self.__weights_init__()
# def __weights_init__(self):
# nn.init.xavier_uniform_(self.conv.weight)
def forward(self, input):
h = self.upsampling(input)
h = self.conv(h)
return h
components/DeConv_Depthwise_ECA.py
+48
View File
@@ -0,0 +1,48 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: DeConv copy.py
# Created Date: Tuesday July 20th 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Saturday, 19th February 2022 6:16:08 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
from tokenize import group
from torch import nn
import math
class DeConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size = 3, upsampl_scale = 2, padding="zero", up_mode = "bilinear"):
super().__init__()
if up_mode.lower() == "bilinear":
self.upsampling = nn.UpsamplingBilinear2d(scale_factor=upsampl_scale)
elif up_mode.lower() == "nearest":
self.upsampling = nn.UpsamplingNearest2d(scale_factor=upsampl_scale)
b = 1
gamma = 2
k_size = int(abs(math.log(out_channels,2)+b)/gamma)
k_size = k_size if k_size % 2 else k_size+1
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.se = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
self.sigmoid = nn.Sigmoid()
padding_size = int((kernel_size -1)/2)
self.conv1x1 = nn.Conv2d(in_channels = in_channels, out_channels = out_channels, kernel_size= 1)
self.conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=padding_size, bias=False, groups=out_channels)
# nn.init.xavier_uniform_(self.conv.weight)
# self.__weights_init__()
# def __weights_init__(self):
# nn.init.xavier_uniform_(self.conv.weight)
def forward(self, input):
h = self.conv1x1(input)
h = self.upsampling(h)
y = self.avg_pool(h)
y = self.se(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
y = self.sigmoid(y)
h = self.conv(h)
return h * y.expand_as(h)
components/ECA.py
+48
View File
@@ -0,0 +1,48 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: ECA.py
# Created Date: Tuesday February 23rd 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Tuesday, 23rd February 2021 9:14:28 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
import math
import torch
from torch import nn
from torch.nn.parameter import Parameter
class eca_layer(nn.Module):
"""Constructs a ECA module.
Args:
channel: Number of channels of the input feature map
k_size: Adaptive selection of kernel size
"""
def __init__(self, channel):
super(eca_layer, self).__init__()
b = 1
gamma = 2
k_size = int(abs(math.log(channel,2)+b)/gamma)
k_size = k_size if k_size % 2 else k_size+1
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# x: input features with shape [b, c, h, w]
# b, c, h, w = x.size()
# feature descriptor on the global spatial information
y = self.avg_pool(x)
# Two different branches of ECA module
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
# Multi-scale information fusion
y = self.sigmoid(y)
return x * y.expand_as(x)
components/ECA_Depthwise_Conv.py
+41
View File
@@ -0,0 +1,41 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: DeConv copy.py
# Created Date: Tuesday July 20th 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Saturday, 19th February 2022 6:15:53 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
from torch import nn
import math
class ECADW(nn.Module):
def __init__(self, in_channels, kernel_size = 3, stride = 2, padding="zero"):
super().__init__()
b = 1
gamma = 2
k_size = int(abs(math.log(in_channels,2)+b)/gamma)
k_size = k_size if k_size % 2 else k_size+1
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.se = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
self.sigmoid = nn.Sigmoid()
padding_size = int((kernel_size -1)/2)
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3,
padding=padding_size, bias=False, groups=in_channels, stride=stride)
# nn.init.xavier_uniform_(self.conv.weight)
# self.__weights_init__()
# def __weights_init__(self):
# nn.init.xavier_uniform_(self.conv.weight)
def forward(self, input):
y = self.avg_pool(input)
y = self.se(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
y = self.sigmoid(y)
h = self.conv(input)
return h * y.expand_as(h)
components/Generator_modulation_depthwise.py → components/Generator_eca_depthwise.py
+72 -48
View File
@@ -5,14 +5,14 @@
# Created Date: Sunday January 16th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Tuesday, 15th February 2022 1:54:50 am
# Last Modified: Saturday, 19th February 2022 6:25:38 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
import torch
from torch import nn
from components.DeConv_Depthwise import DeConv
# from components.DeConv_Invo import DeConv
class Demodule(nn.Module):
@@ -29,21 +29,6 @@ class Demodule(nn.Module):
tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
return x * tmp
class ApplyStyle(nn.Module):
"""
@ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
"""
def __init__(self, latent_size, channels):
super(ApplyStyle, self).__init__()
self.linear = nn.Linear(latent_size, channels * 2)
def forward(self, x, latent):
style = self.linear(latent) # style => [batch_size, n_channels*2]
shape = [-1, 2, x.size(1), 1, 1]
style = style.view(shape) # [batch_size, 2, n_channels, ...]
#x = x * (style[:, 0] + 1.) + style[:, 1]
x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1
return x
class Modulation(nn.Module):
def __init__(self, latent_size, channels):
@@ -59,7 +44,7 @@ class Modulation(nn.Module):
return x
class ResnetBlock_Modulation(nn.Module):
def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True),res_mode="depthwise"):
super(ResnetBlock_Modulation, self).__init__()
p = 0
@@ -72,7 +57,17 @@ class ResnetBlock_Modulation(nn.Module):
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), Demodule()]
if res_mode.lower() == "conv":
conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), Demodule()]
elif res_mode.lower() == "depthwise":
conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "depthwise_eca":
from components.ECA_Depthwise_Conv import ECADW
conv1 += [ECADW(dim, kernel_size=3, padding=p, stride=2),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
self.conv1 = nn.Sequential(*conv1)
self.style1 = Modulation(latent_size, dim)
self.act1 = activation
@@ -87,17 +82,30 @@ class ResnetBlock_Modulation(nn.Module):
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), Demodule()]
if res_mode.lower() == "conv":
conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), Demodule()]
elif res_mode.lower() == "depthwise":
conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "depthwise_eca":
from components.ECA_Depthwise_Conv import ECADW
conv2 += [ECADW(dim, kernel_size=3, padding=p, stride=2),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
self.conv2 = nn.Sequential(*conv2)
self.style2 = Modulation(latent_size, dim)
def forward(self, x, dlatents_in_slice):
y = self.conv1(x)
y = self.style1(y, dlatents_in_slice)
y = self.style1(x, dlatents_in_slice)
y = self.conv1(y)
y = self.act1(y)
y = self.conv2(y)
y = self.style2(y, dlatents_in_slice)
y = self.conv2(y)
out = x + y
return out
@@ -108,68 +116,84 @@ class Generator(nn.Module):
):
super().__init__()
chn = kwargs["g_conv_dim"]
id_dim = kwargs["id_dim"]
k_size = kwargs["g_kernel_size"]
res_num = kwargs["res_num"]
in_channel = kwargs["in_channel"]
up_mode = kwargs["up_mode"]
res_mode = kwargs["res_mode"]
conv_mode = kwargs["conv_mode"]
padding_size= int((k_size -1)/2)
padding_type= 'reflect'
activation = nn.ReLU(True)
from components.ECA_Depthwise_Conv import ECADW
# self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(3, 64, kernel_size=7, padding=0, bias=False),
# nn.BatchNorm2d(64), activation)
self.first_layer = nn.Sequential(nn.ReflectionPad2d(1),
nn.Conv2d(3, 64, kernel_size=3, padding=0, bias=False),
nn.BatchNorm2d(64), activation)
nn.Conv2d(3, in_channel, kernel_size=3, padding=0, bias=False),
nn.BatchNorm2d(in_channel), activation)
# self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(64), activation)
### downsample
self.down1 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, groups=64, padding=1, stride=2),
nn.Conv2d(64, 128, kernel_size=1, bias=False),
nn.BatchNorm2d(128), activation)
self.down1 = nn.Sequential(ECADW(in_channel,kernel_size=3, padding=1, stride=2),
nn.Conv2d(in_channel, in_channel*2, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*2), activation)
self.down2 = nn.Sequential(nn.Conv2d(128, 128, kernel_size=3, groups=128, padding=1, stride=2),
nn.Conv2d(128, 256, kernel_size=1, bias=False),
nn.BatchNorm2d(256), activation)
self.down2 = nn.Sequential(ECADW(in_channel*2, kernel_size=3, padding=1, stride=2),
nn.Conv2d(in_channel*2, in_channel*4, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*4), activation)
self.down3 = nn.Sequential(nn.Conv2d(256, 256, kernel_size=3, groups=256, padding=1, stride=2),
nn.Conv2d(256, 512, kernel_size=1, bias=False),
nn.BatchNorm2d(512), activation)
self.down3 = nn.Sequential(ECADW(in_channel*4, kernel_size=3, padding=1, stride=2),
nn.Conv2d(in_channel*4, in_channel*8, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*8), activation)
self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, groups=512, padding=1, stride=2),
nn.Conv2d(512, 512, kernel_size=1, bias=False),
nn.BatchNorm2d(512), activation)
self.down4 = nn.Sequential(ECADW(in_channel*8, kernel_size=3, padding=1, stride=2),
nn.Conv2d(in_channel*8, in_channel*8, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*8), activation)
### resnet blocks
BN = []
for i in range(res_num):
BN += [
ResnetBlock_Modulation(512, latent_size=chn, padding_type=padding_type, activation=activation)]
ResnetBlock_Modulation(in_channel*8, latent_size=id_dim, padding_type=padding_type, activation=activation)]
self.BottleNeck = nn.Sequential(*BN)
if conv_mode.lower() == "conv":
from components.DeConv import DeConv
Deconv = DeConv
elif conv_mode.lower() == "depthwise":
from components.DeConv_Depthwise import DeConv
Deconv = DeConv
elif conv_mode.lower() == "depthwise_eca":
from components.DeConv_Depthwise_ECA import DeConv
Deconv = DeConv
self.up4 = nn.Sequential(
DeConv(512,512,3),
nn.BatchNorm2d(512), activation
DeConv(in_channel*8,in_channel*8,3),
nn.BatchNorm2d(in_channel*8), activation
)
self.up3 = nn.Sequential(
DeConv(512,256,3),
nn.BatchNorm2d(256), activation
DeConv(in_channel*8,in_channel*4,3),
nn.BatchNorm2d(in_channel*4), activation
)
self.up2 = nn.Sequential(
DeConv(256,128,3),
nn.BatchNorm2d(128), activation
DeConv(in_channel*4,in_channel*2,3),
nn.BatchNorm2d(in_channel*2), activation
)
self.up1 = nn.Sequential(
DeConv(128,64,3),
nn.BatchNorm2d(64), activation
DeConv(in_channel*2,in_channel,3),
nn.BatchNorm2d(in_channel), activation
)
# self.last_layer = nn.Sequential(nn.Conv2d(64, 3, kernel_size=3, padding=1))
self.last_layer = nn.Sequential(nn.ReflectionPad2d(1),
nn.Conv2d(64, 3, kernel_size=3, padding=0))
nn.Conv2d(in_channel, 3, kernel_size=3, padding=0))
# self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
# nn.Conv2d(64, 3, kernel_size=7, padding=0))
components/Generator_modulation_depthwise_config.py
+9 -1
View File
@@ -5,7 +5,7 @@
# Created Date: Sunday January 16th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Thursday, 17th February 2022 2:06:09 am
# Last Modified: Saturday, 19th February 2022 5:16:02 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
@@ -79,6 +79,10 @@ class ResnetBlock_Modulation(nn.Module):
conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "depthwise_eca":
conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
self.conv1 = nn.Sequential(*conv1)
self.style1 = Modulation(latent_size, dim)
self.act1 = activation
@@ -99,6 +103,10 @@ class ResnetBlock_Modulation(nn.Module):
conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "depthwise_eca":
conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
self.conv2 = nn.Sequential(*conv2)
self.style2 = Modulation(latent_size, dim)
speed_test.py
+5 -3
View File
@@ -5,7 +5,7 @@
# Created Date: Thursday February 10th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Thursday, 17th February 2022 2:33:30 am
# Last Modified: Saturday, 19th February 2022 6:23:43 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
@@ -21,7 +21,8 @@ if __name__ == '__main__':
# cudnn.benchmark = True
# cudnn.enabled = True
# script = "Generator_modulation_up"
script = "Generator_modulation_depthwise_config"
script = "Generator_modulation_depthwise"
# script = "Generator_modulation_depthwise_config"
# script = "Generator_ori_config"
class_name = "Generator"
arcface_ckpt= "arcface_ckpt/arcface_checkpoint.tar"
@@ -32,7 +33,8 @@ if __name__ == '__main__':
"res_num": 9,
# "up_mode": "nearest",
"up_mode": "bilinear",
"res_mode": "depthwise"
"res_mode": "depthwise_eca",
"conv_mode": "depthwise_eca"
}
os.environ['CUDA_VISIBLE_DEVICES'] = str(0)