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
)