# -*- coding: utf-8 -*- import os, sys import torch import glob import time, shutil import math import gc from tqdm import tqdm from collections import defaultdict # torch module import from torch.multiprocessing import Pool, Process, set_start_method from torch.utils.tensorboard import SummaryWriter from torch.utils.data import DataLoader try: set_start_method('spawn') except RuntimeError: pass # import files from local folder root_path = os.path.abspath('.') sys.path.append(root_path) from loss.gan_loss import GANLoss, MultiScaleGANLoss from loss.pixel_loss import PixelLoss, L1_Charbonnier_loss from loss.perceptual_loss import PerceptualLoss from loss.anime_perceptual_loss import Anime_PerceptualLoss from architecture.dataset import ImageDataset from scripts.generate_lr_esr import generate_low_res_esr # Mixed precision training scaler = torch.cuda.amp.GradScaler() class train_master(object): def __init__(self, options, args, model_name, has_discriminator=False) -> None: # General specs setup self.args = args self.model_name = model_name self.options = options self.has_discriminator = has_discriminator # Loss init self.loss_init() # Generator self.call_model() # generator + discriminator... # Optimizer self.learning_rate = options['start_learning_rate'] self.optimizer_g = torch.optim.Adam(self.generator.parameters(), lr=self.learning_rate, betas=(options["adam_beta1"], options["adam_beta2"])) if self.has_discriminator: self.optimizer_d = torch.optim.Adam(self.discriminator.parameters(), lr=self.learning_rate, betas=(self.options["adam_beta1"], self.options["adam_beta2"])) # Train specs self.start_iteration = 0 self.lowest_generator_loss = float("inf") # Other auxiliary function self.writer = SummaryWriter() self.weight_store = defaultdict(int) # Options setting self.n_iterations = options['train_iterations'] self.batch_size = options['train_batch_size'] self.n_cpu = options['train_dataloader_workers'] def adjust_learning_rate(self, iteration_idx): self.learning_rate = self.options['start_learning_rate'] end_iteration = self.options['train_iterations'] # Calculate a learning rate we need in real-time based on the iteration_idx for idx in range(min(end_iteration, iteration_idx)//self.options['decay_iteration']): idx = idx+1 if idx * self.options['decay_iteration'] in self.options['double_milestones']: # double the learning rate in milestones self.learning_rate = self.learning_rate * 2 else: # else, try to multiply decay_gamma (when we decay, we won't upscale) self.learning_rate = self.learning_rate * self.options['decay_gamma'] # should be divisible in all cases # Change the learning rate to our target for param_group in self.optimizer_g.param_groups: param_group['lr'] = self.learning_rate if self.has_discriminator: # print("We didn't yet handle discriminator, but we think that it should be necessary") for param_group in self.optimizer_d.param_groups: param_group['lr'] = self.learning_rate assert(self.learning_rate == self.optimizer_g.param_groups[0]['lr']) def pixel_loss_load(self): if self.options['pixel_loss'] == "L1": self.cri_pix = PixelLoss().cuda() elif self.options['pixel_loss'] == "L1_Charbonnier": self.cri_pix = L1_Charbonnier_loss().cuda() print("We are using {} loss".format(self.options['pixel_loss'])) def GAN_loss_load(self): # parameter init gan_loss_weight = self.options["gan_loss_weight"] vgg_type = self.options['train_perceptual_vgg_type'] # Preceptual Loss self.cri_pix = torch.nn.L1Loss().cuda() self.cri_vgg_perceptual = PerceptualLoss(self.options['train_perceptual_layer_weights'], vgg_type, perceptual_weight=self.options["vgg_perceptual_loss_weight"]).cuda() self.cri_danbooru_perceptual = Anime_PerceptualLoss(self.options["Danbooru_layer_weights"], perceptual_weight=self.options["danbooru_perceptual_loss_weight"]).cuda() # GAN loss if self.options['discriminator_type'] == "PatchDiscriminator": self.cri_gan = MultiScaleGANLoss(gan_type="lsgan", loss_weight=gan_loss_weight).cuda() # already put in loss scaler for discriminator elif self.options['discriminator_type'] == "UNetDiscriminator": self.cri_gan = GANLoss(gan_type="vanilla", loss_weight=gan_loss_weight).cuda() # already put in loss scaler for discriminator def tensorboard_epoch_draw(self, epoch_loss, epoch): self.writer.add_scalar('Loss/train-Loss-Epoch', epoch_loss, epoch) def master_run(self): torch.backends.cudnn.benchmark = True print("options are ", self.options) # Generate a new LR dataset before doing anything (Must before Data Loading) self.generate_lr() # Load data train_lr_paths = glob.glob(self.options["lr_dataset_path"] + "/*.*") degrade_hr_paths = glob.glob(self.options["degrade_hr_dataset_path"] + "/*.*") train_hr_paths = glob.glob(self.options["train_hr_dataset_path"] + "/*.*") train_dataloader = DataLoader(ImageDataset(train_lr_paths, degrade_hr_paths, train_hr_paths), batch_size=self.batch_size, shuffle=True, num_workers=self.n_cpu) # ONLY LOAD HALF OF CPU AVAILABLE dataset_length = len(os.listdir(self.options["train_hr_dataset_path"])) # Check if we need to load weight if self.args.auto_resume_best or self.args.auto_resume_closest: self.load_weight(self.model_name) elif self.args.pretrained_path != "": # If we give a pretrained path, we will use it (Should have in GAN training which uses pretrained L1 loss Network) self.load_pretrained(self.model_name) # Start iterating the epochs start_epoch = self.start_iteration // math.ceil(dataset_length / self.options['train_batch_size']) n_epochs = self.n_iterations // math.ceil(dataset_length / self.options['train_batch_size']) iteration_idx = self.start_iteration # init the iteration index self.batch_idx = iteration_idx self.adjust_learning_rate(iteration_idx) # adjust the learning rate to the desired one at the beginning for epoch in range(start_epoch, n_epochs): print("This is epoch {} and the start iteration is {} with learning rate {}".format(epoch, iteration_idx, self.optimizer_g.param_groups[0]['lr'])) # Generate new lr degradation image if epoch != start_epoch and epoch % self.options['degradate_generation_freq'] == 0: self.generate_lr() # Batch training loss_per_epoch = 0.0 self.generator.train() tqdm_bar = tqdm(train_dataloader, total=len(train_dataloader)) for batch_idx, imgs in enumerate(tqdm_bar): imgs_lr = imgs["lr"].cuda() imgs_degrade_hr = imgs["degrade_hr"].cuda() imgs_hr = imgs["hr"].cuda() # Used for each iteration self.generator_loss = 0 self.single_iteration(imgs_lr, imgs_degrade_hr, imgs_hr) # tensorboard and updates self.tensorboard_report(iteration_idx) loss_per_epoch += self.generator_loss.item() ################################# Save model weights and update hyperparameter ######################################## if self.lowest_generator_loss >= self.generator_loss.item(): self.lowest_generator_loss = self.generator_loss.item() print("\nSave model with the lowest generator_loss among all iteartions ", self.lowest_generator_loss) # Store the best self.save_weight(iteration_idx, self.model_name+"_best", self.options) self.lowest_tensorboard_report(iteration_idx) # Update iteration and learning rate iteration_idx += 1 self.batch_idx = iteration_idx if iteration_idx % self.options['decay_iteration'] == 0: self.adjust_learning_rate(iteration_idx) # adjust the learning rate to the desired one print("Update the learning rate to {} at iteration {} ".format(self.optimizer_g.param_groups[0]['lr'], iteration_idx)) # Don't clean any memory here, it will dramatically slow down the code # Per epoch report self.tensorboard_epoch_draw( loss_per_epoch/batch_idx, epoch) # Per epoch store weight self.save_weight(iteration_idx, self.model_name+"_closest", self.options) # Backup Checkpoint (Per 50 epoch) if epoch % self.options['checkpoints_freq'] == 0 or epoch == n_epochs-1: self.save_weight(iteration_idx, "checkpoints/" + self.model_name + "_epoch_" + str(epoch), self.options) # Clean unneeded GPU cache (since we use subprocess for generate_lr(), so we need to kill them all) torch.cuda.empty_cache() time.sleep(5) # For enough time to clean the cache def single_iteration(self, imgs_lr, imgs_degrade_hr, imgs_hr): ############################################# Generator section ################################################## self.optimizer_g.zero_grad() if self.has_discriminator: for p in self.discriminator.parameters(): p.requires_grad = False with torch.cuda.amp.autocast(): # generate high res image gen_hr = self.generator(imgs_lr) # all distinct loss will be stored in self.weight_store (per iteration) self.calculate_loss(gen_hr, imgs_hr) # backward needed loss # self.loss_generator_total.backward() # self.optimizer_g.step() scaler.scale(self.generator_loss).backward() # loss backward scaler.step(self.optimizer_g) scaler.update() ################################################################################################################### if self.has_discriminator: ##################################### Discriminator section ##################################################### for p in self.discriminator.parameters(): p.requires_grad = True self.optimizer_d.zero_grad() # discriminator real input with torch.cuda.amp.autocast(): # We only need imgs_degrade_hr instead of imgs_hr in discriminator (Thus, we don't want to introduce usm in the discriminator) real_d_preds = self.discriminator(imgs_degrade_hr) l_d_real = self.cri_gan(real_d_preds, True, is_disc=True) scaler.scale(l_d_real).backward() # discriminator fake input with torch.cuda.amp.autocast(): fake_d_preds = self.discriminator(gen_hr.detach().clone()) l_d_fake = self.cri_gan(fake_d_preds, False, is_disc=True) scaler.scale(l_d_fake).backward() # update scaler.step(self.optimizer_d) scaler.update() ################################################################################################################## def load_pretrained(self, name): # This part will load generator weight here, and it doesn't need to weight_dir = self.args.pretrained_path if not os.path.exists(weight_dir): print("No such pretrained "+weight_dir+" file exists! We end the program! Please check the dir!") os._exit(0) checkpoint_g = torch.load(weight_dir) if 'model_state_dict' in checkpoint_g: self.generator.load_state_dict(checkpoint_g['model_state_dict']) elif 'params_ema' in checkpoint_g: self.generator.load_state_dict(checkpoint_g['params_ema']) else: raise NotImplementedError("We didn't cannot locate the weight of thie pretrained weight") print(f"We will use pretrained "+name+" weight!") def load_weight(self, head_prefix): # Resume best or the closest weight available head = head_prefix+"_best" if self.args.auto_resume_best else head_prefix+"_closest" if os.path.exists("saved_models/"+head+"_generator.pth"): print("We need to resume previous " + head + " weight") # Generator checkpoint_g = torch.load("saved_models/"+head+"_generator.pth") self.generator.load_state_dict(checkpoint_g['model_state_dict']) self.optimizer_g.load_state_dict(checkpoint_g['optimizer_state_dict']) # Discriminator if self.has_discriminator: checkpoint_d = torch.load("saved_models/"+head+"_discriminator.pth") self.discriminator.load_state_dict(checkpoint_d['model_state_dict']) self.optimizer_d.load_state_dict(checkpoint_d['optimizer_state_dict']) assert(checkpoint_g['iteration'] == checkpoint_d['iteration']) # must be the same for iteration in generator and discriminator self.start_iteration = checkpoint_g['iteration'] + 1 # Prepare lowest generator if os.path.exists("saved_models/" + head_prefix + "_best_generator.pth"): checkpoint_g = torch.load("saved_models/" + head_prefix + "_best_generator.pth") # load generator weight else: print("There is no best weight exists!") self.lowest_generator_loss = min(self.lowest_generator_loss, checkpoint_g["lowest_generator_weight"] ) print("The lowest generator loss at the beginning is ", self.lowest_generator_loss) else: print(f"No saved_models/"+head+"_generator.pth " or " saved_models/"+head+"_discriminator.pth exists") print(f"We will start from the iteration {self.start_iteration}") def save_weight(self, iteration, name, opt): # Generator torch.save({ 'iteration': iteration, 'model_state_dict': self.generator.state_dict(), 'optimizer_state_dict': self.optimizer_g.state_dict(), 'lowest_generator_weight': self.lowest_generator_loss, 'opt': opt, }, "saved_models/" + name + "_generator.pth") # 'pixel_loss': self.weight_store["pixel_loss"], # 'perceptual_loss': self.weight_store['perceptual_loss'], # 'gan_loss': self.weight_store["gan_loss"], if self.has_discriminator: # Discriminator torch.save({ 'iteration': iteration, 'model_state_dict': self.discriminator.state_dict(), 'optimizer_state_dict': self.optimizer_d.state_dict(), }, "saved_models/" + name + "_discriminator.pth") def lowest_tensorboard_report(self, iteration): self.writer.add_scalar('Loss/lowest-weight', self.generator_loss, iteration) @torch.no_grad() def generate_lr(self): # If we directly use API, pytorch2.0 may raise an unknown bugs which is extremely slow on degradation pipeline os.system("python scripts/generate_lr_esr.py") # Assert check lr_paths = os.listdir(self.options["lr_dataset_path"]) degrade_hr_paths = os.listdir(self.options["degrade_hr_dataset_path"]) hr_paths = os.listdir(self.options["train_hr_dataset_path"]) assert(len(lr_paths) == len(degrade_hr_paths)) assert(len(lr_paths) == len(hr_paths))