diff --git a/GUI/file_sync/filestate_machine0.json b/GUI/file_sync/filestate_machine0.json index 6079e3f..260ae6c 100644 --- a/GUI/file_sync/filestate_machine0.json +++ b/GUI/file_sync/filestate_machine0.json @@ -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 } \ No newline at end of file diff --git a/components/Generator_Res_config.py b/components/Generator_Res_config.py index 65c8a95..fa41c4a 100644 --- a/components/Generator_Res_config.py +++ b/components/Generator_Res_config.py @@ -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), diff --git a/components/Generator_Res_config1.py b/components/Generator_Res_config1.py new file mode 100644 index 0000000..b5674ae --- /dev/null +++ b/components/Generator_Res_config1.py @@ -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 \ No newline at end of file diff --git a/start_train.sh b/start_train.sh index 8de53eb..20bd78c 100644 --- a/start_train.sh +++ b/start_train.sh @@ -1,3 +1,3 @@ -nohup python train_multigpu.py > cycle_lstu1.log 2>&1 & \ No newline at end of file +nohup python train_multigpu.py > cycle_res1.log 2>&1 & \ No newline at end of file diff --git a/train_yamls/train_cycleloss_res.yaml b/train_yamls/train_cycleloss_res.yaml index ac9832e..181d367 100644 --- a/train_yamls/train_cycleloss_res.yaml +++ b/train_yamls/train_cycleloss_res.yaml @@ -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