From 11a55c9165843cd7168ff934eac537a6706102a2 Mon Sep 17 00:00:00 2001 From: chenxuanhong Date: Tue, 15 Feb 2022 12:16:29 +0800 Subject: [PATCH] update --- components/Generator_oriae_modulation.py | 198 +++++++++++++++++++++++ train_multigpu.py | 12 +- train_yamls/train_oriae_modulation.yaml | 63 ++++++++ 3 files changed, 267 insertions(+), 6 deletions(-) create mode 100644 components/Generator_oriae_modulation.py create mode 100644 train_yamls/train_oriae_modulation.yaml diff --git a/components/Generator_oriae_modulation.py b/components/Generator_oriae_modulation.py new file mode 100644 index 0000000..feb4057 --- /dev/null +++ b/components/Generator_oriae_modulation.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- +############################################################# +# File: Generator.py +# Created Date: Sunday January 16th 2022 +# Author: Chen Xuanhong +# Email: chenxuanhongzju@outlook.com +# Last Modified: Tuesday, 15th February 2022 12:03:17 pm +# Modified By: Chen Xuanhong +# Copyright (c) 2022 Shanghai Jiao Tong University +############################################################# + +from audioop import bias +import torch +from torch import nn +from torch.nn import init +from torch.nn import functional as F + +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)): + 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) + conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), 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) + conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), 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.act1(y) + y = self.conv2(y) + y = self.style2(y, dlatents_in_slice) + out = x + y + return out + + +class Generator(nn.Module): + def __init__( + self, + **kwargs + ): + super().__init__() + + chn = kwargs["g_conv_dim"] + k_size = kwargs["g_kernel_size"] + res_num = kwargs["res_num"] + + padding_size= int((k_size -1)/2) + padding_type= 'reflect' + + activation = nn.ReLU(True) + + self.first_layer = nn.Sequential(nn.ReflectionPad2d(1), nn.Conv2d(3, 64, kernel_size=3, padding=0, bias=False), + nn.BatchNorm2d(64), activation) + ### downsample + self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(128), activation) + + self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(256), activation) + + self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(512), activation) + + self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(512), activation) + + ### resnet blocks + BN = [] + for i in range(res_num): + BN += [ + ResnetBlock_Modulation(512, latent_size=chn, padding_type=padding_type, activation=activation)] + self.BottleNeck = nn.Sequential(*BN) + + self.up4 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(512), activation + ) + + self.up3 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(256), activation + ) + + self.up2 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(128), activation + ) + + self.up1 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(64), activation + ) + + self.last_layer = nn.Sequential(nn.ReflectionPad2d(1), nn.Conv2d(64, 3, kernel_size=3, 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 diff --git a/train_multigpu.py b/train_multigpu.py index fd98c92..1f2b62f 100644 --- a/train_multigpu.py +++ b/train_multigpu.py @@ -5,7 +5,7 @@ # Created Date: Tuesday April 28th 2020 # Author: Chen Xuanhong # Email: chenxuanhongzju@outlook.com -# Last Modified: Monday, 14th February 2022 11:54:02 pm +# Last Modified: Tuesday, 15th February 2022 12:06:30 pm # Modified By: Chen Xuanhong # Copyright (c) 2020 Shanghai Jiao Tong University ############################################################# @@ -31,9 +31,9 @@ def getParameters(): parser = argparse.ArgumentParser() # general settings - parser.add_argument('-v', '--version', type=str, default='depthwise', + parser.add_argument('-v', '--version', type=str, default='oriae_modulation', help="version name for train, test, finetune") - parser.add_argument('-t', '--tag', type=str, default='depthwise_conv', + parser.add_argument('-t', '--tag', type=str, default='oriae_modulation', help="tag for current experiment") parser.add_argument('-p', '--phase', type=str, default="train", @@ -46,13 +46,13 @@ def getParameters(): # training parser.add_argument('--experiment_description', type=str, - default="使用depthwise卷积作为基础算子测试性能") + default="验证是否是Decoder导致的发紫") - parser.add_argument('--train_yaml', type=str, default="train_Depthwise.yaml") + parser.add_argument('--train_yaml', type=str, default="train_oriae_modulation.yaml") # system logger parser.add_argument('--logger', type=str, - default="wandb", choices=['tensorboard', 'wandb','none'], help='system logger') + default="tensorboard", choices=['tensorboard', 'wandb','none'], help='system logger') # # logs (does not to be changed in most time) # parser.add_argument('--dataloader_workers', type=int, default=6) diff --git a/train_yamls/train_oriae_modulation.yaml b/train_yamls/train_oriae_modulation.yaml new file mode 100644 index 0000000..2f24fc3 --- /dev/null +++ b/train_yamls/train_oriae_modulation.yaml @@ -0,0 +1,63 @@ +# Related scripts +train_script_name: multi_gpu + +# models' scripts +model_configs: + g_model: + script: Generator_oriae_modulation + class_name: Generator + module_params: + g_conv_dim: 512 + g_kernel_size: 3 + res_num: 9 + + d_model: + script: projected_discriminator + class_name: ProjectedDiscriminator + module_params: + diffaug: False + interp224: False + backbone_kwargs: {} + +arcface_ckpt: arcface_ckpt/arcface_checkpoint.tar + +# Training information +batch_size: 64 + +# Dataset +dataloader: VGGFace2HQ_multigpu +dataset_name: vggface2_hq +dataset_params: + random_seed: 1234 + dataloader_workers: 4 + +eval_dataloader: DIV2K_hdf5 +eval_dataset_name: DF2K_H5_Eval +eval_batch_size: 2 + +# Dataset + +# Optimizer +optim_type: Adam +g_optim_config: + lr: 0.0004 + betas: [ 0, 0.99] + eps: !!float 1e-8 + +d_optim_config: + lr: 0.0004 + betas: [ 0, 0.99] + eps: !!float 1e-8 + +id_weight: 20.0 +reconstruct_weight: 10.0 +feature_match_weight: 10.0 + +# Log +log_step: 300 +model_save_step: 10000 +sample_step: 1000 +total_step: 1000000 +checkpoint_names: + generator_name: Generator + discriminator_name: Discriminator \ No newline at end of file