From db049166a04523a85b67b4bf58b04b037f368d9b Mon Sep 17 00:00:00 2001 From: chenxuanhong Date: Sat, 19 Feb 2022 18:26:22 +0800 Subject: [PATCH] eca depth wise --- components/DeConv.py | 41 ++++++ components/DeConv_Depthwise_ECA.py | 48 +++++++ components/ECA.py | 48 +++++++ components/ECA_Depthwise_Conv.py | 41 ++++++ ...epthwise.py => Generator_eca_depthwise.py} | 120 +++++++++++------- .../Generator_modulation_depthwise_config.py | 10 +- speed_test.py | 8 +- 7 files changed, 264 insertions(+), 52 deletions(-) create mode 100644 components/DeConv.py create mode 100644 components/DeConv_Depthwise_ECA.py create mode 100644 components/ECA.py create mode 100644 components/ECA_Depthwise_Conv.py rename components/{Generator_modulation_depthwise.py => Generator_eca_depthwise.py} (56%) diff --git a/components/DeConv.py b/components/DeConv.py new file mode 100644 index 0000000..e5d7af9 --- /dev/null +++ b/components/DeConv.py @@ -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 \ No newline at end of file diff --git a/components/DeConv_Depthwise_ECA.py b/components/DeConv_Depthwise_ECA.py new file mode 100644 index 0000000..c06755b --- /dev/null +++ b/components/DeConv_Depthwise_ECA.py @@ -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) \ No newline at end of file diff --git a/components/ECA.py b/components/ECA.py new file mode 100644 index 0000000..595e730 --- /dev/null +++ b/components/ECA.py @@ -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) \ No newline at end of file diff --git a/components/ECA_Depthwise_Conv.py b/components/ECA_Depthwise_Conv.py new file mode 100644 index 0000000..b9c21e3 --- /dev/null +++ b/components/ECA_Depthwise_Conv.py @@ -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) \ No newline at end of file diff --git a/components/Generator_modulation_depthwise.py b/components/Generator_eca_depthwise.py similarity index 56% rename from components/Generator_modulation_depthwise.py rename to components/Generator_eca_depthwise.py index f632d4f..b922c2a 100644 --- a/components/Generator_modulation_depthwise.py +++ b/components/Generator_eca_depthwise.py @@ -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)) diff --git a/components/Generator_modulation_depthwise_config.py b/components/Generator_modulation_depthwise_config.py index 8348b8b..2e03dc5 100644 --- a/components/Generator_modulation_depthwise_config.py +++ b/components/Generator_modulation_depthwise_config.py @@ -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) diff --git a/speed_test.py b/speed_test.py index b952b48..2a3eb84 100644 --- a/speed_test.py +++ b/speed_test.py @@ -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)