Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import pytorch_lightning as pl | |
| import torch | |
| from skimage.transform import resize | |
| from lib.common.train_util import batch_mean | |
| from lib.net import NormalNet | |
| class Normal(pl.LightningModule): | |
| def __init__(self, cfg): | |
| super(Normal, self).__init__() | |
| self.cfg = cfg | |
| self.batch_size = self.cfg.batch_size | |
| self.lr_F = self.cfg.lr_netF | |
| self.lr_B = self.cfg.lr_netB | |
| self.lr_D = self.cfg.lr_netD | |
| self.overfit = cfg.overfit | |
| self.F_losses = [item[0] for item in self.cfg.net.front_losses] | |
| self.B_losses = [item[0] for item in self.cfg.net.back_losses] | |
| self.ALL_losses = self.F_losses + self.B_losses | |
| self.automatic_optimization = False | |
| self.schedulers = [] | |
| self.netG = NormalNet(self.cfg) | |
| self.in_nml = [item[0] for item in cfg.net.in_nml] | |
| # Training related | |
| def configure_optimizers(self): | |
| optim_params_N_D = None | |
| optimizer_N_D = None | |
| scheduler_N_D = None | |
| # set optimizer | |
| optim_params_N_F = [{"params": self.netG.netF.parameters(), "lr": self.lr_F}] | |
| optim_params_N_B = [{"params": self.netG.netB.parameters(), "lr": self.lr_B}] | |
| optimizer_N_F = torch.optim.Adam(optim_params_N_F, lr=self.lr_F, betas=(0.5, 0.999)) | |
| optimizer_N_B = torch.optim.Adam(optim_params_N_B, lr=self.lr_B, betas=(0.5, 0.999)) | |
| scheduler_N_F = torch.optim.lr_scheduler.MultiStepLR( | |
| optimizer_N_F, milestones=self.cfg.schedule, gamma=self.cfg.gamma | |
| ) | |
| scheduler_N_B = torch.optim.lr_scheduler.MultiStepLR( | |
| optimizer_N_B, milestones=self.cfg.schedule, gamma=self.cfg.gamma | |
| ) | |
| if 'gan' in self.ALL_losses: | |
| optim_params_N_D = [{"params": self.netG.netD.parameters(), "lr": self.lr_D}] | |
| optimizer_N_D = torch.optim.Adam(optim_params_N_D, lr=self.lr_D, betas=(0.5, 0.999)) | |
| scheduler_N_D = torch.optim.lr_scheduler.MultiStepLR( | |
| optimizer_N_D, milestones=self.cfg.schedule, gamma=self.cfg.gamma | |
| ) | |
| self.schedulers = [scheduler_N_F, scheduler_N_B, scheduler_N_D] | |
| optims = [optimizer_N_F, optimizer_N_B, optimizer_N_D] | |
| else: | |
| self.schedulers = [scheduler_N_F, scheduler_N_B] | |
| optims = [optimizer_N_F, optimizer_N_B] | |
| return optims, self.schedulers | |
| def render_func(self, render_tensor, dataset, idx): | |
| height = render_tensor["image"].shape[2] | |
| result_list = [] | |
| for name in render_tensor.keys(): | |
| result_list.append( | |
| resize( | |
| ((render_tensor[name].cpu().numpy()[0] + 1.0) / 2.0).transpose(1, 2, 0), | |
| (height, height), | |
| anti_aliasing=True, | |
| ) | |
| ) | |
| self.logger.log_image( | |
| key=f"Normal/{dataset}/{idx if not self.overfit else 1}", | |
| images=[(np.concatenate(result_list, axis=1) * 255.0).astype(np.uint8)] | |
| ) | |
| def training_step(self, batch, batch_idx): | |
| self.netG.train() | |
| # retrieve the data | |
| in_tensor = {} | |
| for name in self.in_nml: | |
| in_tensor[name] = batch[name] | |
| FB_tensor = {"normal_F": batch["normal_F"], "normal_B": batch["normal_B"]} | |
| in_tensor.update(FB_tensor) | |
| preds_F, preds_B = self.netG(in_tensor) | |
| error_dict = self.netG.get_norm_error(preds_F, preds_B, FB_tensor) | |
| if 'gan' in self.ALL_losses: | |
| (opt_F, opt_B, opt_D) = self.optimizers() | |
| opt_F.zero_grad() | |
| self.manual_backward(error_dict["netF"]) | |
| opt_B.zero_grad() | |
| self.manual_backward(error_dict["netB"], retain_graph=True) | |
| opt_D.zero_grad() | |
| self.manual_backward(error_dict["netD"]) | |
| opt_F.step() | |
| opt_B.step() | |
| opt_D.step() | |
| else: | |
| (opt_F, opt_B) = self.optimizers() | |
| opt_F.zero_grad() | |
| self.manual_backward(error_dict["netF"]) | |
| opt_B.zero_grad() | |
| self.manual_backward(error_dict["netB"]) | |
| opt_F.step() | |
| opt_B.step() | |
| if batch_idx > 0 and batch_idx % int( | |
| self.cfg.freq_show_train | |
| ) == 0 and self.cfg.devices == 1: | |
| self.netG.eval() | |
| with torch.no_grad(): | |
| nmlF, nmlB = self.netG(in_tensor) | |
| in_tensor.update({"nmlF": nmlF, "nmlB": nmlB}) | |
| self.render_func(in_tensor, "train", self.global_step) | |
| # metrics processing | |
| metrics_log = {"loss": error_dict["netF"] + error_dict["netB"]} | |
| if "gan" in self.ALL_losses: | |
| metrics_log["loss"] += error_dict["netD"] | |
| for key in error_dict.keys(): | |
| metrics_log["train/loss_" + key] = error_dict[key].item() | |
| self.log_dict( | |
| metrics_log, prog_bar=True, logger=True, on_step=True, on_epoch=False, sync_dist=True | |
| ) | |
| return metrics_log | |
| def training_epoch_end(self, outputs): | |
| # metrics processing | |
| metrics_log = {} | |
| for key in outputs[0].keys(): | |
| if "/" in key: | |
| [stage, loss_name] = key.split("/") | |
| else: | |
| stage = "train" | |
| loss_name = key | |
| metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key) | |
| self.log_dict( | |
| metrics_log, | |
| prog_bar=False, | |
| logger=True, | |
| on_step=False, | |
| on_epoch=True, | |
| rank_zero_only=True | |
| ) | |
| def validation_step(self, batch, batch_idx): | |
| self.netG.eval() | |
| self.netG.training = False | |
| # retrieve the data | |
| in_tensor = {} | |
| for name in self.in_nml: | |
| in_tensor[name] = batch[name] | |
| FB_tensor = {"normal_F": batch["normal_F"], "normal_B": batch["normal_B"]} | |
| in_tensor.update(FB_tensor) | |
| preds_F, preds_B = self.netG(in_tensor) | |
| error_dict = self.netG.get_norm_error(preds_F, preds_B, FB_tensor) | |
| if batch_idx % int(self.cfg.freq_show_train) == 0 and self.cfg.devices == 1: | |
| with torch.no_grad(): | |
| nmlF, nmlB = self.netG(in_tensor) | |
| in_tensor.update({"nmlF": nmlF, "nmlB": nmlB}) | |
| self.render_func(in_tensor, "val", batch_idx) | |
| # metrics processing | |
| metrics_log = {"val/loss": error_dict["netF"] + error_dict["netB"]} | |
| if "gan" in self.ALL_losses: | |
| metrics_log["val/loss"] += error_dict["netD"] | |
| for key in error_dict.keys(): | |
| metrics_log["val/" + key] = error_dict[key].item() | |
| return metrics_log | |
| def validation_epoch_end(self, outputs): | |
| # metrics processing | |
| metrics_log = {} | |
| for key in outputs[0].keys(): | |
| [stage, loss_name] = key.split("/") | |
| metrics_log[f"{stage}/avg-{loss_name}"] = batch_mean(outputs, key) | |
| self.log_dict( | |
| metrics_log, | |
| prog_bar=False, | |
| logger=True, | |
| on_step=False, | |
| on_epoch=True, | |
| rank_zero_only=True | |
| ) | |