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-27 20:23:37 +08:00
parent dc11678ed6
commit 860713323d
6 changed files with 289 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
breakpoint.json
+1 -1
View File
@@ -1,6 +1,6 @@
{
"breakpoint": [
1053,
1054,
0
]
}
check_list.txt
+7
View File
@@ -869,3 +869,10 @@ n001111\0426_02.jpg
n001112\0164_01.jpg
n001112\0186_02.jpg
n001112\0234_01.jpg
n001113\0190_01.jpg
n001113\0289_01.jpg
n001113\0290_01.jpg
n001113\0293_01.jpg
n001113\0302_01.jpg
n001113\0423_01.jpg
n001113\0443_02.jpg
components/Generator_Invobn_config2.py
+1 -1
View File
@@ -5,7 +5,7 @@
# Created Date: Saturday February 26th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Sunday, 27th February 2022 3:58:23 pm
# Last Modified: Sunday, 27th February 2022 7:50:18 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
components/Generator_Invobn_config3.py
+274
View File
@@ -0,0 +1,274 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: Generator_Invobn_config1.py
# Created Date: Saturday February 26th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Sunday, 27th February 2022 7:50:18 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
import torch
from torch import nn
# from components.DeConv_Invo import DeConv
class Demodule(nn.Module):
def __init__(self, epsilon=1e-8):
"""
@notice: avoid in-place ops.
https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
"""
super(Demodule, self).__init__()
self.epsilon = epsilon
def forward(self, x):
tmp = torch.mul(x, x) # or x ** 2
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):
super(Modulation, self).__init__()
self.linear = nn.Linear(latent_size, channels)
def forward(self, x, latent):
style = self.linear(latent) # style => [batch_size, n_channels*2]
shape = [-1, x.size(1), 1, 1]
style = style.view(shape) # [batch_size, 2, n_channels, ...]
#x = x * (style[:, 0] + 1.) + style[:, 1]
x = x * style
return x
class ResnetBlock_Modulation(nn.Module):
def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True),res_mode="depthwise"):
super(ResnetBlock_Modulation, self).__init__()
p = 0
conv1 = []
if padding_type == 'reflect':
conv1 += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv1 += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
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":
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() == "invo":
from components.misc.Involution_BN import involution
conv1 += [involution(dim,3, 1, 0),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "eca_invo":
from components.misc.Involution_ECA import involution
conv1 += [involution(dim, 3, 1, 0),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
self.conv1 = nn.Sequential(*conv1)
self.style1 = Modulation(latent_size, dim)
self.act1 = activation
p = 0
conv2 = []
if padding_type == 'reflect':
conv2 += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv2 += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
res_mode = "conv"
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":
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() == "invo":
from components.misc.Involution_BN import involution
conv2 += [involution(dim, 3, 1,0),
nn.Conv2d(dim, dim, kernel_size=1),
Demodule()]
elif res_mode.lower() == "eca_invo":
from components.misc.Involution_ECA import involution
conv2 += [involution(dim, 3, 1, 0),
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.style1(x, dlatents_in_slice)
y = self.conv1(y)
y = self.act1(y)
y = self.style2(y, dlatents_in_slice)
y = self.conv2(y)
out = x + y
return out
class Generator(nn.Module):
def __init__(
self,
**kwargs
):
super().__init__()
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"]
aggregator = kwargs["aggregator"]
res_mode = aggregator
padding_size= int((k_size -1)/2)
padding_type= 'reflect'
activation = nn.ReLU(True)
if aggregator == "invo":
from components.misc.Involution_BN import involution
from components.DeConv_Invobn import DeConv
elif aggregator == "eca_invo":
from components.misc.Involution_ECA import involution
from components.DeConv_ECA_Invo import DeConv
# 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)
# 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(
involution(in_channel,3,2),
nn.Conv2d(in_channel, in_channel*2, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*2),
activation)
self.down2 = nn.Sequential(
involution(in_channel*2,3,2),
nn.Conv2d(in_channel*2, in_channel*4, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*4),
activation)
self.down3 = nn.Sequential(
involution(in_channel*4,3,2),
nn.Conv2d(in_channel*4, in_channel*8, kernel_size=1, bias=False),
nn.BatchNorm2d(in_channel*8),
activation)
self.down4 = nn.Sequential(
involution(in_channel*8,3,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(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,up_mode=up_mode),
nn.BatchNorm2d(in_channel*8),
activation
)
self.up3 = nn.Sequential(
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,up_mode=up_mode),
nn.BatchNorm2d(in_channel*2),
activation
)
self.up1 = nn.Sequential(
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),
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))
# self.__weights_init__()
# def __weights_init__(self):
# for layer in self.encoder:
# if isinstance(layer,nn.Conv2d):
# nn.init.xavier_uniform_(layer.weight)
# for layer in self.encoder2:
# if isinstance(layer,nn.Conv2d):
# nn.init.xavier_uniform_(layer.weight)
def forward(self, img, id):
res = self.first_layer(img)
res = self.down1(res)
res = self.down2(res)
res = self.down3(res)
res = self.down4(res)
for i in range(len(self.BottleNeck)):
res = self.BottleNeck[i](res, id)
res = self.up4(res)
res = self.up3(res)
res = self.up2(res)
res = self.up1(res)
res = self.last_layer(res)
return res
flops.py
+5 -4
View File
@@ -5,7 +5,7 @@
# Created Date: Sunday February 13th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Sunday, 27th February 2022 6:30:55 pm
# Last Modified: Sunday, 27th February 2022 8:15:11 pm
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
@@ -21,8 +21,9 @@ from thop import clever_format
if __name__ == '__main__':
#
# script = "Generator_modulation_up"
script = "Generator_Invobn_config2"
# script = "Generator_modulation_up"
script = "Generator_Invobn_config3"
# script = "Generator_ori_config"
# script = "Generator_ori_config"
class_name = "Generator"
arcface_ckpt= "arcface_ckpt/arcface_checkpoint.tar"
@@ -30,7 +31,7 @@ if __name__ == '__main__':
"id_dim": 512,
"g_kernel_size": 3,
"in_channel":16,
"res_num": 6,
"res_num": 9,
# "up_mode": "nearest",
"up_mode": "bilinear",
"aggregator": "eca_invo"
train_yamls/train_Invobn_config.yaml
+1 -1
View File
@@ -4,7 +4,7 @@ train_script_name: multi_gpu
# models' scripts
model_configs:
g_model:
script: Generator_Invobn_config
script: Generator_Invobn_config3
class_name: Generator
module_params:
id_dim: 512