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

update

Browse Source
This commit is contained in:
chenxuanhong
2022-02-17 16:30:33 +08:00
parent d82e46f0e7
commit 913e4916d4
17 changed files with 45302 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/Generator_modulation_depthwise_config.py
+48 -22
View File
@@ -5,11 +5,12 @@
# Created Date: Sunday January 16th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Tuesday, 15th February 2022 12:52:22 pm
# Last Modified: Thursday, 17th February 2022 2:06:09 am
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
from audioop import bias
import torch
from torch import nn
from components.DeConv_Depthwise import DeConv
@@ -59,7 +60,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 +73,12 @@ 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()]
self.conv1 = nn.Sequential(*conv1)
self.style1 = Modulation(latent_size, dim)
self.act1 = activation
@@ -87,17 +93,25 @@ 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()]
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
@@ -112,6 +126,8 @@ class Generator(nn.Module):
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"]
padding_size= int((k_size -1)/2)
padding_type= 'reflect'
@@ -121,52 +137,62 @@ class Generator(nn.Module):
# 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, in_channel, kernel_size=3, padding=0, bias=False),
nn.BatchNorm2d(in_channel), 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(in_channel, in_channel, kernel_size=3, groups=in_channel, padding=1, stride=2),
nn.Conv2d(in_channel, in_channel*2, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*2), activation)
nn.BatchNorm2d(in_channel*2),
activation)
self.down2 = nn.Sequential(nn.Conv2d(in_channel*2, in_channel*2, kernel_size=3, groups=in_channel*2, padding=1, stride=2),
nn.Conv2d(in_channel*2, in_channel*4, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*4), activation)
nn.BatchNorm2d(in_channel*4),
activation)
self.down3 = nn.Sequential(nn.Conv2d(in_channel*4, in_channel*4, kernel_size=3, groups=in_channel*4, padding=1, stride=2),
nn.Conv2d(in_channel*4, in_channel*8, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*8), activation)
nn.BatchNorm2d(in_channel*8),
activation)
self.down4 = nn.Sequential(nn.Conv2d(in_channel*8, in_channel*8, kernel_size=3, groups=in_channel*8, padding=1, stride=2),
nn.Conv2d(in_channel*8, in_channel*8, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*8), activation)
nn.BatchNorm2d(in_channel*8),
activation)
### resnet blocks
BN = []
for i in range(res_num):
BN += [
ResnetBlock_Modulation(in_channel*8, latent_size=id_dim, padding_type=padding_type, activation=activation)]
ResnetBlock_Modulation(in_channel*8, latent_size=id_dim,
padding_type=padding_type, activation=activation, res_mode=res_mode)]
self.BottleNeck = nn.Sequential(*BN)
self.up4 = nn.Sequential(
DeConv(in_channel*8,in_channel*8,3),
nn.BatchNorm2d(in_channel*8), activation
DeConv(in_channel*8,in_channel*8,3,up_mode=up_mode),
nn.BatchNorm2d(in_channel*8),
activation
)
self.up3 = nn.Sequential(
DeConv(in_channel*8,in_channel*4,3),
nn.BatchNorm2d(in_channel*4), activation
DeConv(in_channel*8,in_channel*4,3,up_mode=up_mode),
nn.BatchNorm2d(in_channel*4),
activation
)
self.up2 = nn.Sequential(
DeConv(in_channel*4,in_channel*2,3),
nn.BatchNorm2d(in_channel*2), activation
DeConv(in_channel*4,in_channel*2,3,up_mode=up_mode),
nn.BatchNorm2d(in_channel*2),
activation
)
self.up1 = nn.Sequential(
DeConv(in_channel*2,in_channel,3),
nn.BatchNorm2d(in_channel), activation
DeConv(in_channel*2,in_channel,3,up_mode=up_mode),
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),