This commit is contained in:
XHChen0528
2022-03-24 00:32:14 +08:00
parent 34c80c5315
commit fe52d5fbd5
5 changed files with 383 additions and 17 deletions
+5 -3
View File
@@ -1,6 +1,6 @@
{
"GUI.py": 1647657822.9152665,
"test.py": 1647657822.945273,
"test.py": 1647879709.2723496,
"train.py": 1647657822.9562755,
"components\\Generator.py": 1647657822.93127,
"components\\projected_discriminator.py": 1647657822.938272,
@@ -60,7 +60,7 @@
"face_crop.py": 1647657822.9422722,
"face_crop_video.py": 1647657822.9422722,
"similarity.py": 1647657822.945273,
"train_multigpu.py": 1647700474.445049,
"train_multigpu.py": 1647967698.603863,
"components\\arcface_decoder.py": 1647657822.9352713,
"components\\Generator_nobias.py": 1643179001.810856,
"data_tools\\data_loader_VGGFace2HQ_multigpu.py": 1647657822.9402719,
@@ -186,5 +186,7 @@
"train_scripts\\trainer_distillation_mgpu_withrec_importweight.py": 1647657822.9592762,
"train_scripts\\trainer_multi_gpu_CUT.py": 1647676964.475,
"train_scripts\\trainer_multi_gpu_cycle.py": 1647705628.7020626,
"components\\Generator_LSTU_config.py": 1647704615.1532204
"components\\Generator_LSTU_config.py": 1647954099.1135788,
"components\\Generator_Res_config.py": 1648006159.4385264,
"train_yamls\\train_cycleloss_res.yaml": 1648006232.5734456
}
+11 -12
View File
@@ -120,14 +120,14 @@ class ResUpSampleBlock(nn.Module):
latent_size,
activation=nn.LeakyReLU(0.2),
res_mode="depthwise"):
super(ResnetBlock_Adain, self).__init__()
super(ResUpSampleBlock, self).__init__()
conv1 = []
self.in1 = InstanceNorm()
self.in2 = InstanceNorm()
if res_mode.lower() == "conv":
conv1 += [activation,
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)]
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv1 += [activation,
@@ -145,7 +145,7 @@ class ResUpSampleBlock(nn.Module):
conv2 = []
if res_mode.lower() == "conv":
conv2 += [activation,
nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)]
nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv2 += [activation,
@@ -183,14 +183,14 @@ class ResDownSampleBlock(nn.Module):
out_dim,
activation=nn.LeakyReLU(0.2),
res_mode="depthwise"):
super(ResnetBlock_Adain, self).__init__()
super(ResDownSampleBlock, self).__init__()
conv1 = []
if res_mode.lower() == "conv":
conv1 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1)]
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv1 += [
@@ -213,7 +213,7 @@ class ResDownSampleBlock(nn.Module):
conv2 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)]
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv2 += [
@@ -235,7 +235,7 @@ class ResDownSampleBlock(nn.Module):
def forward(self, x):
y = self.conv1(y)
y = self.conv1(x)
y = self.resampling(y)
y = self.conv2(y)
res = self.reshape1_1(x)
@@ -264,7 +264,6 @@ class Generator(nn.Module):
padding_type= 'reflect'
activation = nn.LeakyReLU(0.2)
from components.DeConv_Depthwise 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)
@@ -310,7 +309,7 @@ class Generator(nn.Module):
# padding_type=padding_type, activation=activation, res_mode=res_mode)]
# self.BottleNeck = nn.Sequential(*BN)
self.up4 = ResDownSampleBlock(in_channel*8,in_channel*8,id_dim,res_mode=res_mode) # 64
self.up4 = ResUpSampleBlock(in_channel*8,in_channel*8,id_dim,res_mode=res_mode) # 64
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*8, in_channel*8, kernel_size=3, stride=1, padding=1, bias=False),
@@ -318,7 +317,7 @@ class Generator(nn.Module):
# activation
# )
self.up3 = ResDownSampleBlock(in_channel*8,in_channel*4,id_dim,res_mode=res_mode) # 128
self.up3 = ResUpSampleBlock(in_channel*8,in_channel*4,id_dim,res_mode=res_mode) # 128
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*8, in_channel*4, kernel_size=3, stride=1, padding=1, bias=False),
@@ -326,7 +325,7 @@ class Generator(nn.Module):
# activation
# )
self.up2 = ResDownSampleBlock(in_channel*4,in_channel*2,id_dim,res_mode=res_mode) # 256
self.up2 = ResUpSampleBlock(in_channel*4,in_channel*2,id_dim,res_mode=res_mode) # 256
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*4, in_channel*2, kernel_size=3, stride=1, padding=1, bias=False),
@@ -334,7 +333,7 @@ class Generator(nn.Module):
# activation
# )
self.up1 = ResDownSampleBlock(in_channel*2,in_channel,id_dim,res_mode=res_mode) # 512
self.up1 = ResUpSampleBlock(in_channel*2,in_channel,id_dim,res_mode=res_mode) # 512
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+365
View File
@@ -0,0 +1,365 @@
#!/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.LSTU import LSTU
# from components.DeConv_Invo import DeConv
class InstanceNorm(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(InstanceNorm, self).__init__()
self.epsilon = epsilon
def forward(self, x):
x = x - torch.mean(x, (2, 3), True)
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 ResnetBlock_Adain(nn.Module):
def __init__(self,
dim,
latent_size,
activation=nn.LeakyReLU(0.2),
res_mode="depthwise"):
super(ResnetBlock_Adain, self).__init__()
conv1 = []
self.in1 = InstanceNorm()
self.in2 = InstanceNorm()
if res_mode.lower() == "conv":
conv1 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv1 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv1 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1)]
self.conv1 = nn.Sequential(*conv1)
self.style1 = ApplyStyle(latent_size, dim)
conv2 = []
if res_mode.lower() == "conv":
conv2 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv2 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv2 += [activation,
nn.Conv2d(dim, dim, kernel_size=3, padding=1,groups=dim, bias=False),
nn.Conv2d(dim, dim, kernel_size=1)]
self.conv2 = nn.Sequential(*conv2)
self.style2 = ApplyStyle(latent_size, dim)
def forward(self, x, dlatents_in_slice):
y = self.in1(x)
y = self.style1(y, dlatents_in_slice)
y = self.conv1(y)
y = self.in2(y)
y = self.style2(y, dlatents_in_slice)
y = self.conv2(y)
out = x + y
return out
class ResUpSampleBlock(nn.Module):
def __init__(self,
in_dim,
out_dim,
latent_size,
activation=nn.LeakyReLU(0.2),
res_mode="depthwise"):
super(ResUpSampleBlock, self).__init__()
conv1 = []
self.in1 = InstanceNorm()
self.in2 = InstanceNorm()
if res_mode.lower() == "conv":
conv1 += [activation,
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv1 += [activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=in_dim, bias=False),
nn.Conv2d(in_dim, out_dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv1 += [activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=in_dim, bias=False),
nn.Conv2d(in_dim, out_dim, kernel_size=1)]
self.conv1 = nn.Sequential(*conv1)
self.style1 = ApplyStyle(latent_size, in_dim)
conv2 = []
if res_mode.lower() == "conv":
conv2 += [activation,
nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv2 += [activation,
nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1,groups=out_dim, bias=False),
nn.Conv2d(out_dim, out_dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv2 += [activation,
nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1,groups=out_dim, bias=False),
nn.Conv2d(out_dim, out_dim, kernel_size=1)]
self.conv2 = nn.Sequential(*conv2)
self.style2 = ApplyStyle(latent_size, out_dim)
self.reshape1_1 = nn.Conv2d(in_dim, out_dim, kernel_size=1)
self.resampling = nn.UpsamplingBilinear2d(scale_factor=2)
def forward(self, x, dlatents_in_slice):
y = self.in1(x)
y = self.style1(y, dlatents_in_slice)
y = self.conv1(y)
y = self.resampling(y)
y = self.in2(y)
y = self.style2(y, dlatents_in_slice)
y = self.conv2(y)
res = self.reshape1_1(x)
res = self.resampling(res)
out = res + y
return out
class ResDownSampleBlock(nn.Module):
def __init__(self,
in_dim,
out_dim,
activation=nn.LeakyReLU(0.2),
res_mode="depthwise"):
super(ResDownSampleBlock, self).__init__()
conv1 = []
if res_mode.lower() == "conv":
conv1 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv1 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=in_dim, bias=False),
nn.Conv2d(in_dim, in_dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv1 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=in_dim, bias=False),
nn.Conv2d(in_dim, in_dim, kernel_size=1)]
self.conv1 = nn.Sequential(*conv1)
conv2 = []
if res_mode.lower() == "conv":
conv2 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1, bias=False)]
elif res_mode.lower() == "depthwise":
conv2 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=out_dim, bias=False),
nn.Conv2d(in_dim, out_dim, kernel_size=1)]
elif res_mode.lower() == "depthwise_eca":
conv2 += [
nn.BatchNorm2d(in_dim),
activation,
nn.Conv2d(in_dim, in_dim, kernel_size=3, padding=1,groups=out_dim, bias=False),
nn.Conv2d(in_dim, out_dim, kernel_size=1)]
self.conv2 = nn.Sequential(*conv2)
self.reshape1_1 = nn.Conv2d(in_dim, out_dim, kernel_size=1)
self.resampling = nn.AvgPool2d(3,2,1)
def forward(self, x):
y = self.conv1(x)
y = self.resampling(y)
y = self.conv2(y)
res = self.reshape1_1(x)
res = self.resampling(res)
out = res + 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 = kwargs["res_mode"]
padding_size= int((k_size -1)/2)
padding_type= 'reflect'
activation = nn.LeakyReLU(0.2)
# 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, stride=2, padding=0, bias=False),
nn.BatchNorm2d(in_channel),
activation) # 256
# self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(64), activation)
### downsample
self.down1 = ResDownSampleBlock(in_channel, in_channel*2,res_mode=res_mode)
# nn.Sequential(
# nn.Conv2d(in_channel, in_channel*2, stride=2, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*2),
# activation) # 128
self.down2 = ResDownSampleBlock(in_channel*2, in_channel*4,res_mode=res_mode)
# nn.Sequential(
# nn.Conv2d(in_channel*2, in_channel*4, stride=2, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*4),
# activation) # 64
# self.lstu = LSTU(in_channel*4,in_channel*4,in_channel*8,4)
self.down3 = ResDownSampleBlock(in_channel*4, in_channel*8,res_mode=res_mode)
# nn.Sequential(
# nn.Conv2d(in_channel*4, in_channel*8, stride=2, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*8),
# activation) # 32
# self.down4 = nn.Sequential(
# nn.Conv2d(in_channel*8, in_channel*8, stride=2, kernel_size=3, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*8),
# activation)
### resnet blocks
BN = []
for i in range(res_num):
BN += [
ResnetBlock_Adain(in_channel*8, latent_size=id_dim,res_mode=res_mode)]
self.BottleNeck = nn.Sequential(*BN)
self.up4 = ResUpSampleBlock(in_channel*8,in_channel*8,id_dim,res_mode=res_mode) # 64
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*8, in_channel*8, kernel_size=3, stride=1, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*8),
# activation
# )
self.up3 = ResUpSampleBlock(in_channel*8,in_channel*4,id_dim,res_mode=res_mode) # 128
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*8, in_channel*4, kernel_size=3, stride=1, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*4),
# activation
# )
self.up2 = ResUpSampleBlock(in_channel*4,in_channel*2,id_dim,res_mode=res_mode) # 256
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*4, in_channel*2, kernel_size=3, stride=1, padding=1, bias=False),
# nn.BatchNorm2d(in_channel*2),
# activation
# )
self.up1 = ResUpSampleBlock(in_channel*2,in_channel,id_dim,res_mode=res_mode) # 512
# nn.Sequential(
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
# 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)
for i in range(len(self.BottleNeck)):
res = self.BottleNeck[i](res, id)
res = self.up4(res,id)
res = self.up3(res,id)
res = self.up2(res,id) # + skip
res = self.up1(res,id)
res = self.last_layer(res)
return res
+1 -1
View File
@@ -1,3 +1,3 @@
nohup python train_multigpu.py > cycle_lstu1.log 2>&1 &
nohup python train_multigpu.py > cycle_res1.log 2>&1 &
+1 -1
View File
@@ -27,7 +27,7 @@ model_configs:
arcface_ckpt: arcface_ckpt/arcface_checkpoint.tar
# Training information
batch_size: 20
batch_size: 18
# Dataset
dataloader: VGGFace2HQ_multigpu