import numpy as np import torch import torch.nn.functional as F import torch.distributed as dist import os import random import copy import cv2 import math import lpips from utils.commons.hparams import hparams from utils.commons.tensor_utils import tensors_to_scalars, convert_to_np, move_to_cuda from utils.nn.model_utils import not_requires_grad, num_params from utils.commons.dataset_utils import data_loader from utils.nn.schedulers import NoneSchedule from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint, restore_weights, restore_opt_state from tasks.os_avatar.loss_utils.vgg19_loss import VGG19Loss from tasks.os_avatar.dataset_utils.motion2video_dataset import Motion2Video_Dataset from tasks.os_avatar.img2plane_task import OSAvatarImg2PlaneTask from modules.eg3ds.models.triplane import TriPlaneGenerator from modules.eg3ds.models.dual_discriminator import DualDiscriminator, SingleDiscriminator from modules.eg3ds.torch_utils.ops import conv2d_gradfix from modules.eg3ds.torch_utils.ops import upfirdn2d from modules.eg3ds.models.dual_discriminator import filtered_resizing from modules.real3d.secc_img2plane import OSAvatarSECC_Img2plane from deep_3drecon.secc_renderer import SECC_Renderer from data_util.face3d_helper import Face3DHelper from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter from data_gen.runs.binarizer_nerf import get_lip_rect from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor from inference.edit_secc import blink_eye_for_secc class ScheduleForLM3DImg2PlaneEG3D(NoneSchedule): def __init__(self, optimizer, lr, lr_d, warmup_updates=0): self.optimizer = optimizer self.constant_lr = self.lr = lr self.lr_d = lr_d self.warmup_updates = warmup_updates self.step(0) def step(self, num_updates): constant_lr = self.constant_lr if self.warmup_updates > 0 and num_updates <= self.warmup_updates: warmup = min(num_updates / self.warmup_updates, 1.0) self.lr = max(constant_lr * warmup, 1e-7) else: self.lr = constant_lr for optim_i in range(len(self.optimizer)-1): lr_mul_cano_img2plane = hparams['lr_mul_cano_img2plane'] * min(1.0, num_updates / (hparams['start_adv_iters']+20000)) self.optimizer[optim_i].param_groups[0]['lr'] = lr_mul_cano_img2plane * self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000)) if num_updates > 6_000 else 0 # cano_img2plane self.optimizer[optim_i].param_groups[0]['lr'] = max(5e-6, self.optimizer[optim_i].param_groups[0]['lr']) if num_updates >= hparams['stop_update_i2p_iters']: self.optimizer[optim_i].param_groups[0]['lr'] = 0. self.optimizer[optim_i].param_groups[1]['lr'] = max(5e-6, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) if num_updates > 0 else 0 # secc_img2plane self.optimizer[optim_i].param_groups[2]['lr'] = max(5e-6, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) if num_updates > 6_000 else 0 # decoder self.optimizer[optim_i].param_groups[3]['lr'] = max(5e-6, self.lr * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) if num_updates > 30_000 else 0 # sr self.optimizer[-1].param_groups[0]['lr'] = max(5e-6, self.lr_d * (hparams.get("lr_decay_rate", 0.95)) ** (num_updates // hparams.get("lr_decay_interval", 5_000))) # for disc return self.lr class SECC_Img2PlaneEG3DTask(OSAvatarImg2PlaneTask): def __init__(self): super().__init__() self.seg_model = MediapipeSegmenter() self.dataset_cls = Motion2Video_Dataset self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68') def build_model(self): self.eg3d_model = TriPlaneGenerator() load_ckpt(self.eg3d_model, hparams['pretrained_eg3d_ckpt'], strict=True) self.model = OSAvatarSECC_Img2plane() self.disc = DualDiscriminator() self.cano_img2plane_params = [p for k, p in self.model.cano_img2plane_backbone.named_parameters() if p.requires_grad] self.secc_img2plane_params = [p for k, p in self.model.secc_img2plane_backbone.named_parameters() if p.requires_grad] self.decoder_params = [p for p in self.model.decoder.parameters() if p.requires_grad] self.upsample_params = [p for p in self.model.superresolution.parameters() if p.requires_grad] self.disc_params = [p for k, p in self.disc.named_parameters() if p.requires_grad] if hparams.get("add_ffhq_singe_disc", False): self.ffhq_disc = DualDiscriminator() self.disc_params += [p for k, p in self.ffhq_disc.named_parameters() if p.requires_grad] eg3d_dir = 'checkpoints/geneface2_ckpts/eg3d_baseline_run2' load_ckpt(self.ffhq_disc, eg3d_dir, model_name='disc', strict=True) self.secc_renderer = SECC_Renderer(512) if hparams.get('init_from_ckpt', '') != '': ckpt_dir = hparams['init_from_ckpt'] try: load_ckpt(self.model.cano_img2plane_backbone, ckpt_dir, model_name='model.cano_img2plane_backbone', strict=True) load_ckpt(self.model.secc_img2plane_backbone, ckpt_dir, model_name='model.secc_img2plane_backbone', strict=True) except: # from a img2plane ckpt load_ckpt(self.model.cano_img2plane_backbone, ckpt_dir, model_name='model.img2plane_backbone', strict=False) # load_ckpt(self.model.cano_img2plane_backbone, ckpt_dir, model_name='model.img2plane_backbone', strict=True) load_ckpt(self.model.decoder, ckpt_dir, model_name='model.decoder', strict=True) load_ckpt(self.model.superresolution, ckpt_dir, model_name='model.superresolution', strict=False) # false for spade sr load_ckpt(self.disc, ckpt_dir, model_name='disc', strict=True) return self.model def build_optimizer(self, model): self.optimizer_gen = optimizer_gen = torch.optim.Adam( self.cano_img2plane_params, lr=hparams['lr_g'], # we use a 0.5x smaller lr for transformer betas=(hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g']) ) self.optimizer_gen.add_param_group({ 'params': self.secc_img2plane_params, 'lr': hparams['lr_g'], 'betas': (hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g']) }) self.optimizer_gen.add_param_group({ 'params': self.decoder_params, 'lr': hparams['lr_g'], 'betas': (hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g']) }) self.optimizer_gen.add_param_group({ 'params': self.upsample_params, 'lr': hparams['lr_g'], 'betas': (hparams['optimizer_adam_beta1_g'], hparams['optimizer_adam_beta2_g']) }) mb_ratio_d = hparams['reg_interval_d'] / (hparams['reg_interval_d'] + 1) self.optimizer_disc = optimizer_disc = torch.optim.Adam( self.disc_params, lr=hparams['lr_d'] * mb_ratio_d, betas=(hparams['optimizer_adam_beta1_d'] ** mb_ratio_d, hparams['optimizer_adam_beta2_d'] ** mb_ratio_d)) optimizers = [optimizer_gen, optimizer_disc] return optimizers def build_scheduler(self, optimizer): mb_ratio_d = hparams['reg_interval_d'] / (hparams['reg_interval_d'] + 1) return ScheduleForLM3DImg2PlaneEG3D(optimizer, hparams['lr_g'], hparams['lr_d'] * mb_ratio_d, hparams['warmup_updates']) def forward_G(self, img, camera, cond=None, ret=None, update_emas=False, cache_backbone=True, use_cached_backbone=False): """ ref_img: [B, 3, W, H] camera: [b, 25], 16 dim c2w, and 9 dim intrinsic cond: a dict of cano_secc, tgt_secc, src_secc """ G = self.model gen_output = G.forward(img=img, camera=camera, cond=cond, ret=ret, update_emas=update_emas, cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone) return gen_output def forward_D(self, img, camera, update_emas=False): D = self.disc logits = D.forward(img, camera, update_emas=update_emas) return logits def forward_ffhq_D(self, img, camera, update_emas=False): D = self.ffhq_disc logits = D.forward(img, 0*camera, update_emas=update_emas) return logits def prepare_batch(self, batch): out_batch = {} out_batch['th1kh_ref_cameras'] = batch['th1kh_ref_cameras'] out_batch['th1kh_ref_head_imgs'] = batch['th1kh_ref_head_imgs'] out_batch['th1kh_ref_head_imgs_raw'] = filtered_resizing(batch['th1kh_ref_head_imgs'], size=hparams['neural_rendering_resolution'], f=self.resample_filter, filter_mode='antialiased') out_batch['th1kh_mv_cameras'] = batch['th1kh_mv_cameras'] out_batch['th1kh_mv_head_imgs'] = batch['th1kh_mv_head_imgs'] out_batch['th1kh_mv_head_imgs_raw'] = filtered_resizing(batch['th1kh_mv_head_imgs'], size=hparams['neural_rendering_resolution'], f=self.resample_filter, filter_mode='antialiased') out_batch['th1kh_ref_eulers'] = batch['th1kh_ref_eulers'] out_batch['th1kh_ref_trans'] = batch['th1kh_ref_trans'] with torch.no_grad(): _, out_batch['th1kh_cano_secc'] = self.secc_renderer(batch['th1kh_ref_ids'],batch['th1kh_ref_exps']*0,batch['th1kh_ref_eulers']*0,batch['th1kh_ref_trans']*0) _, out_batch['th1kh_ref_secc'] = self.secc_renderer(batch['th1kh_ref_ids'],batch['th1kh_ref_exps'],batch['th1kh_ref_eulers']*0,batch['th1kh_ref_trans']*0) _, out_batch['th1kh_mv_secc'] = self.secc_renderer(batch['th1kh_mv_ids'],batch['th1kh_mv_exps'],batch['th1kh_mv_eulers']*0,batch['th1kh_mv_trans']*0) if (self.global_step+1) % hparams['reg_interval_g_cond'] == 0: if random.random() < hparams.get("pertube_ref_prob", 0.25): # 1/4的可能对ref secc做pertube out_batch['th1kh_pertube_secc0'] = out_batch['th1kh_ref_secc'].clone() if hparams.get("secc_pertube_mode", 'randn') == 'randn': _, out_batch['th1kh_pertube_secc1'] = self.secc_renderer(batch['th1kh_ref_ids'] + torch.randn_like(batch['th1kh_ref_ids'])*hparams['secc_pertube_randn_scale'],batch['th1kh_ref_exps'] + torch.randn_like(batch['th1kh_ref_exps'])*hparams['secc_pertube_randn_scale'] ,batch['th1kh_ref_eulers']*0,batch['th1kh_ref_trans']*0) elif hparams.get("secc_pertube_mode", 'randn') in ['tv', 'laplacian']: _, out_batch['th1kh_pertube_secc1'] = self.secc_renderer(batch['th1kh_ref_ids'],batch['th1kh_ref_pertube_exps_1'],batch['th1kh_ref_eulers']*0,batch['th1kh_ref_trans']*0) _, out_batch['th1kh_pertube_secc2'] = self.secc_renderer(batch['th1kh_ref_ids'],batch['th1kh_ref_pertube_exps_2'],batch['th1kh_ref_eulers']*0,batch['th1kh_ref_trans']*0) else: raise NotImplementedError() else: out_batch['th1kh_pertube_secc0'] = out_batch['th1kh_mv_secc'] if hparams.get("secc_pertube_mode", 'randn') == 'randn': _, out_batch['th1kh_pertube_secc1'] = self.secc_renderer(batch['th1kh_mv_ids'] + torch.randn_like(batch['th1kh_mv_ids'])*hparams['secc_pertube_randn_scale'],batch['th1kh_mv_exps'] + torch.randn_like(batch['th1kh_mv_exps'])*hparams['secc_pertube_randn_scale'] ,batch['th1kh_mv_eulers']*0,batch['th1kh_mv_trans']*0) elif hparams.get("secc_pertube_mode", 'randn') in ['tv', 'laplacian']: _, out_batch['th1kh_pertube_secc1'] = self.secc_renderer(batch['th1kh_mv_ids'],batch['th1kh_mv_pertube_exps_1'],batch['th1kh_mv_eulers']*0,batch['th1kh_mv_trans']*0) _, out_batch['th1kh_pertube_secc2'] = self.secc_renderer(batch['th1kh_mv_ids'],batch['th1kh_mv_pertube_exps_2'],batch['th1kh_mv_eulers']*0,batch['th1kh_mv_trans']*0) else: raise NotImplementedError() if (self.global_step+1) % hparams['reg_interval_g_cond'] == 0: blink_secc_lst1 = [] blink_secc_lst2 = [] blink_secc_lst3 = [] for i in range(len(out_batch['th1kh_mv_secc'])): if random.random() < 0.25: # 1/4的可能对ref secc做pertube secc = out_batch['th1kh_ref_secc'][i] else: secc = out_batch['th1kh_mv_secc'][i] blink_percent1 = random.random() * 0.5 # 0~0.5 blink_percent3 = 0.5 + random.random() * 0.5 # 0.5~1.0 blink_percent2 = (blink_percent1 + blink_percent3)/2 try: out_secc1 = blink_eye_for_secc(secc, blink_percent1).to(secc.device) out_secc2 = blink_eye_for_secc(secc, blink_percent2).to(secc.device) out_secc3 = blink_eye_for_secc(secc, blink_percent3).to(secc.device) except: print("blink eye for secc failed, use original secc") out_secc1 = copy.deepcopy(secc) out_secc2 = copy.deepcopy(secc) out_secc3 = copy.deepcopy(secc) blink_secc_lst1.append(out_secc1) blink_secc_lst2.append(out_secc2) blink_secc_lst3.append(out_secc3) out_batch['th1kh_blink_mv_secc1'] = torch.stack(blink_secc_lst1) out_batch['th1kh_blink_mv_secc2'] = torch.stack(blink_secc_lst2) out_batch['th1kh_blink_mv_secc3'] = torch.stack(blink_secc_lst3) out_batch['th1kh_ref_head_masks'] = batch['th1kh_ref_head_masks'].unsqueeze(1).bool() out_batch['th1kh_ref_head_masks_raw'] = torch.nn.functional.interpolate(out_batch['th1kh_ref_head_masks'].float(), size=(128,128), mode='nearest').bool() out_batch['th1kh_ref_head_masks_dilate'] = self.dilate_mask(out_batch['th1kh_ref_head_masks'].float(), ksize=41).bool() out_batch['th1kh_ref_head_masks_raw_dilate'] = torch.nn.functional.interpolate(out_batch['th1kh_ref_head_masks_dilate'].float(), size=(128,128), mode='nearest').bool() out_batch['th1kh_mv_head_masks'] = batch['th1kh_mv_head_masks'].unsqueeze(1).bool() out_batch['th1kh_mv_head_masks_raw'] = torch.nn.functional.interpolate(out_batch['th1kh_mv_head_masks'].float(), size=(128,128), mode='nearest').bool() out_batch['th1kh_mv_head_masks_dilate'] = self.dilate_mask(out_batch['th1kh_mv_head_masks'].float(), ksize=41).long() out_batch['th1kh_mv_head_masks_raw_dilate'] = torch.nn.functional.interpolate(out_batch['th1kh_mv_head_masks_dilate'].float(), size=(128,128), mode='nearest').bool() WH = 512 # now we only support 512x512 ref_lm2ds = WH * self.face3d_helper.reconstruct_lm2d(batch['th1kh_ref_ids'],batch['th1kh_ref_exps'],batch['th1kh_ref_eulers'],batch['th1kh_ref_trans']).cpu().numpy() mv_lm2ds = WH * self.face3d_helper.reconstruct_lm2d(batch['th1kh_mv_ids'],batch['th1kh_mv_exps'],batch['th1kh_mv_eulers'],batch['th1kh_mv_trans']).cpu().numpy() ref_lip_rects = [get_lip_rect(ref_lm2ds[i], WH, WH) for i in range(len(ref_lm2ds))] mv_lip_rects = [get_lip_rect(mv_lm2ds[i], WH, WH) for i in range(len(mv_lm2ds))] out_batch['th1kh_ref_lip_rects'] = ref_lip_rects out_batch['th1kh_mv_lip_rects'] = mv_lip_rects return out_batch def run_G_th1kh_src2src_image(self, batch): """ 不在src2src上训练会导致画质变差、不像说话人, 这很合理, 因为i2p也是这样需要update on ref_mse 尤其是在靠近src的画质变好, 但同时会导致depth和color在靠近src的时候闪烁. 解决方法: 算secc2plane pertube loss的时候, 更频繁地在src secc附近计算loss; target到更小的pertube loss """ losses = {} ret = {} ret['losses'] = {} if self.global_step % hparams['update_src2src_interval'] != 0: return losses with torch.autograd.profiler.record_function('G_th1kh_ref_forward'): camera = batch['th1kh_ref_cameras'] img = batch['th1kh_ref_head_imgs'] img_raw = batch['th1kh_ref_head_imgs_raw'] gen_img = self.forward_G(batch['th1kh_ref_head_imgs'], camera, cond={'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_ref_secc'], 'ref_head_img': batch['th1kh_ref_head_imgs'], # used for spade sr 'ref_alphas': batch['th1kh_ref_head_masks'].float(), 'ref_cameras': batch['th1kh_ref_cameras'], }, ret=ret) losses.update(ret['losses']) if hparams.get("masked_error", True): # 之所以用L1不用MSE, 原因是mse对mismatch的pixel loss过大, 而导致面部细节被忽略, 此外还有过模糊的问题 # 对mse raw图像, 因为deform的原因背景没法全黑, 导致这部分mse过高, 我们将其mask掉, 只计算人脸部分 # 在算lpips的时候, 尝试过把非头部mask掉再输入到VGG里面, 但是似乎有点问题, 所以最终没有mask掉非脸 losses['G_th1kh_ref_img_mae_raw'] = self.masked_error_loss(gen_img['image_raw'], img_raw, batch['th1kh_ref_head_masks_raw_dilate'], mode='l1', unmasked_weight=0.2) losses['G_th1kh_ref_img_mae'] = self.masked_error_loss(gen_img['image'], img, batch['th1kh_ref_head_masks_dilate'], mode='l1', unmasked_weight=0.2) pred_img_for_vgg = gen_img['image'] pred_img_raw_for_vgg = gen_img['image_raw'] losses['G_th1kh_ref_img_lpips'] = self.criterion_lpips(pred_img_for_vgg, img).mean() losses['G_th1kh_ref_img_lpips_raw'] = self.criterion_lpips(pred_img_raw_for_vgg, img_raw).mean() disc_inp_img = { 'image': pred_img_for_vgg, 'image_raw': pred_img_raw_for_vgg, } # lip loss batch_size = len(gen_img['image']) lip_mse_loss = 0 lip_lpips_loss = 0 for i in range(batch_size): xmin, xmax, ymin, ymax = batch['th1kh_ref_lip_rects'][i] lip_tgt_imgs = img[i:i+1,:, ymin:ymax,xmin:xmax].contiguous() lip_pred_imgs = pred_img_for_vgg[i:i+1,:, ymin:ymax,xmin:xmax].contiguous() lip_mse_loss = lip_mse_loss + 1/batch_size * (lip_pred_imgs - lip_tgt_imgs).abs().mean() try: lip_lpips_loss = lip_lpips_loss + 1/batch_size * self.criterion_lpips(lip_pred_imgs, lip_tgt_imgs).mean() except: pass losses['G_th1kh_ref_img_lip_mae'] = lip_mse_loss losses['G_th1kh_ref_img_lip_lpips'] = lip_lpips_loss else: losses['G_th1kh_ref_img_mae_raw'] = (gen_img['image_raw'] - img_raw).abs().mean() losses['G_th1kh_ref_img_mae'] = (gen_img['image'] - img).abs().mean() losses['G_th1kh_ref_img_lpips'] = self.criterion_lpips(gen_img['image'], img).mean() losses['G_th1kh_ref_img_lpips_raw'] = self.criterion_lpips(gen_img['image_raw'], img_raw).mean() disc_inp_img = { 'image': gen_img['image'], 'image_raw': gen_img['image_raw'], } # ablate后发现, 去掉对ref的weights reg loss, 会导致学到的density比较散, 略微降低画质 alphas = gen_img['weights_img'].clamp(1e-5, 1 - 1e-5) losses['G_th1kh_ref_weights_entropy_loss'] = torch.mean(- alphas * torch.log2(alphas) - (1 - alphas) * torch.log2(1 - alphas)) face_mask = batch['th1kh_ref_head_masks_raw'].bool() nonface_mask = ~ batch['th1kh_ref_head_masks_raw'].bool() losses['G_th1kh_ref_weights_l1_loss'] = (alphas[nonface_mask]-0).pow(2).mean() + (alphas[face_mask]-1).pow(2).mean() gen_logits = self.forward_D(disc_inp_img, camera) losses['G_th1kh_ref_adv'] = torch.nn.functional.softplus(-gen_logits).mean() return losses def run_G_th1kh_src2tgt_image(self, batch): losses = {} ret = {} ret['losses'] = {} with torch.autograd.profiler.record_function('G_th1kh_mv_forward'): camera = batch['th1kh_mv_cameras'] img = batch['th1kh_mv_head_imgs'] img_raw = batch['th1kh_mv_head_imgs_raw'] gen_img = self.forward_G(batch['th1kh_ref_head_imgs'], camera, cond={'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_mv_secc'], 'ref_head_img': batch['th1kh_ref_head_imgs'], 'ref_alphas': batch['th1kh_ref_head_masks'].float(), 'ref_cameras': batch['th1kh_ref_cameras'], }, ret=ret) losses.update(ret['losses']) self.gen_tmp_output['th1kh_recon_mv_imgs'] = gen_img['image'].detach() self.gen_tmp_output['th1kh_recon_mv_imgs_raw'] = gen_img['image_raw'].detach() losses['G_ref_plane_l1_mean'] = (gen_img['plane'][:,:]).detach().abs().mean() losses['G_ref_plane_l1_std'] = (gen_img['plane'][:,:]).detach().abs().std() if hparams.get("masked_error", True): # 之所以用L1不用MSE, 原因是mse对mismatch的pixel loss过大, 而导致面部细节被忽略, 此外还有过模糊的问题 # 对raw图像, 因为deform的原因背景没法全黑, 导致这部分mse过高, 我们将其mask掉, 只计算人脸部分 losses['G_th1kh_mv_img_mae_raw'] = self.masked_error_loss(gen_img['image_raw'], img_raw, batch['th1kh_mv_head_masks_raw_dilate'], mode='l1', unmasked_weight=0.2) losses['G_th1kh_mv_img_mae'] = self.masked_error_loss(gen_img['image'], img, batch['th1kh_mv_head_masks_dilate'], mode='l1', unmasked_weight=0.2) pred_img_for_vgg = gen_img['image'] pred_img_raw_for_vgg = gen_img['image_raw'] losses['G_th1kh_mv_img_lpips'] = self.criterion_lpips(pred_img_for_vgg, img).mean() losses['G_th1kh_mv_img_lpips_raw'] = self.criterion_lpips(pred_img_raw_for_vgg, img_raw).mean() # emphasize lip loss batch_size = len(gen_img['image']) lip_mse_loss = 0 lip_lpips_loss = 0 for i in range(batch_size): xmin, xmax, ymin, ymax = batch['th1kh_mv_lip_rects'][i] lip_tgt_imgs = img[i:i+1,:, ymin:ymax,xmin:xmax].contiguous() lip_pred_imgs = pred_img_for_vgg[i:i+1,:, ymin:ymax,xmin:xmax].contiguous() lip_mse_loss = lip_mse_loss + 1/batch_size * (lip_pred_imgs - lip_tgt_imgs).abs().mean() try: lip_lpips_loss = lip_lpips_loss + 1/batch_size * self.criterion_lpips(lip_pred_imgs, lip_tgt_imgs).mean() except: pass losses['G_th1kh_mv_img_lip_mae'] = lip_mse_loss losses['G_th1kh_mv_img_lip_lpips'] = lip_lpips_loss self.gen_tmp_output['th1kh_recon_mv_imgs'] = pred_img_for_vgg.detach() self.gen_tmp_output['th1kh_recon_mv_imgs_raw'] = pred_img_raw_for_vgg.detach() disc_inp_img = { 'image': pred_img_for_vgg, 'image_raw': pred_img_raw_for_vgg, } else: losses['G_th1kh_mv_img_mae_raw'] = (gen_img['image_raw'] - img_raw).abs().mean() losses['G_th1kh_mv_img_mae'] = (gen_img['image'] - img).abs().mean() losses['G_th1kh_mv_img_lpips'] = self.criterion_lpips(gen_img['image'], img).mean() losses['G_th1kh_mv_img_lpips_raw'] = self.criterion_lpips(gen_img['image_raw'], img_raw).mean() disc_inp_img = { 'image': gen_img['image'], 'image_raw': gen_img['image_raw'], } alphas = gen_img['weights_img'].clamp(1e-5, 1 - 1e-5) losses['G_th1kh_mv_weights_entropy_loss'] = torch.mean(- alphas * torch.log2(alphas) - (1 - alphas) * torch.log2(1 - alphas)) face_mask = batch['th1kh_mv_head_masks_raw'].bool() nonface_mask = ~ batch['th1kh_mv_head_masks_raw'].bool() losses['G_th1kh_mv_weights_l1_loss'] = (alphas[nonface_mask]-0).pow(2).mean() + (alphas[face_mask]-1).pow(2).mean() gen_logits = self.forward_D(disc_inp_img, camera) losses['G_th1kh_mv_adv'] = torch.nn.functional.softplus(-gen_logits).mean() if hparams.get("add_ffhq_singe_disc", False): gen_logits = self.forward_ffhq_D(disc_inp_img, camera) losses['G_ffhq_adv_maxmimize_model_pred_mv'] = torch.nn.functional.softplus(-gen_logits).mean() return losses def run_G_reg(self, batch): losses = {} imgs = batch['th1kh_ref_head_imgs'] if (self.global_step+1) % hparams['reg_interval_g'] == 0: with torch.autograd.profiler.record_function('G_regularize_forward'): cond={'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_mv_secc'], 'ref_cameras': batch['th1kh_ref_cameras'], 'ref_alphas': batch['th1kh_ref_head_masks'].float(), } initial_coordinates = torch.rand((imgs.shape[0], 1000, 3), device=imgs.device) - 0.5 # [-0.5,0.5] perturbed_coordinates = initial_coordinates + torch.randn_like(initial_coordinates) * 5e-3 all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) source_sigma = self.model.sample(coordinates=all_coordinates, directions=torch.randn_like(all_coordinates), img=imgs, cond=cond, update_emas=False)['sigma'] source_sigma_initial = source_sigma[:, :source_sigma.shape[1]//2] source_sigma_perturbed = source_sigma[:, source_sigma.shape[1]//2:] density_reg_loss = torch.nn.functional.l1_loss(source_sigma_initial, source_sigma_perturbed) # we want the pertubed position has similar density losses['G_th1kh_regularize_density_l1'] = density_reg_loss return losses def run_G_reg_cond(self, batch): losses = {} if (self.global_step+1) % hparams['reg_interval_g_cond'] == 0: # Reg pertube ref/mv_secc, see prepare_batch, we have 25% prob pertube ref and 75% pertube mv. cond = {'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_pertube_secc0']} pertube_cond = {'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_pertube_secc1']} secc_plane = self.model.cal_secc_plane(cond) pertube_secc_plane = self.model.cal_secc_plane(pertube_cond) with torch.autograd.profiler.record_function('G_regularize_forward'): if hparams.get("secc_pertube_mode", 'randn') in ['randn', 'tv']: secc_reg_loss = torch.nn.functional.l1_loss(secc_plane, pertube_secc_plane) # we want the pertubed position has similar density elif hparams.get("secc_pertube_mode", 'randn') == 'laplacian': pertube_cond2 = { 'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_pertube_secc2']} pertube_secc_plane2 = self.model.cal_secc_plane(pertube_cond2) interpolate_secc_plane = (pertube_secc_plane + pertube_secc_plane2) / 2 secc_reg_loss = torch.nn.functional.l1_loss(secc_plane, interpolate_secc_plane) else: raise NotImplementedError() losses['G_th1kh_regularize_pertube_secc_mae'] = secc_reg_loss # Reg blinks blink_cond1 = {'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_blink_mv_secc1']} blink_cond2 = {'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_blink_mv_secc2']} blink_cond3 = {'cond_cano': batch['th1kh_cano_secc'], 'cond_src': batch['th1kh_ref_secc'], 'cond_tgt': batch['th1kh_blink_mv_secc3']} blink_secc_plane1 = self.model.cal_secc_plane(blink_cond1) blink_secc_plane2 = self.model.cal_secc_plane(blink_cond2) blink_secc_plane3 = self.model.cal_secc_plane(blink_cond3) interpolate_blink_secc_plane = (blink_secc_plane1 + blink_secc_plane3) / 2 blink_reg_loss = torch.nn.functional.l1_loss(blink_secc_plane2, interpolate_blink_secc_plane) losses['G_th1kh_regularize_blink_secc_mae'] = blink_reg_loss return losses def forward_D_main(self, batch): """ we update ema this substep. """ FFHQ_DISC_UPDATE_INTERVAL = 4 losses = {} with torch.autograd.profiler.record_function('D_minimize_fake_forward'): # gt ref img & Dmain ref_cameras = batch['th1kh_ref_cameras'] ref_img_tmp_image = batch['th1kh_ref_head_imgs'].detach().requires_grad_(True) ref_img_tmp_image_raw = batch['th1kh_ref_head_imgs_raw'].detach().requires_grad_(True) th1kh_ref_img_tmp = {'image': ref_img_tmp_image, 'image_raw': ref_img_tmp_image_raw} th1kh_ref_logits = self.forward_D(th1kh_ref_img_tmp, ref_cameras) losses['D_th1kh_maximize_gt_ref'] = torch.nn.functional.softplus(-th1kh_ref_logits).mean() if hparams.get("add_ffhq_singe_disc", False) and (self.global_step+1) % FFHQ_DISC_UPDATE_INTERVAL == 0: ffhq_ref_img_tmp = {'image': batch['ffhq_head_imgs'].detach().requires_grad_(True),'image_raw': batch['ffhq_head_imgs_raw'].detach().requires_grad_(True)} ffhq_ref_logits = self.forward_ffhq_D(ffhq_ref_img_tmp, ref_cameras) # ref_camera will be mul 0 in forward_ffhq_D losses['D_ffhq_maximize_gt_ref'] = torch.nn.functional.softplus(-ffhq_ref_logits).mean() # gt ref img & gradient penalty if (self.global_step + 1) % hparams['reg_interval_d'] == 0 and self.training is True: with conv2d_gradfix.no_weight_gradients(): ref_r1_grads = torch.autograd.grad(outputs=[th1kh_ref_logits.sum()], inputs=[th1kh_ref_img_tmp['image'], th1kh_ref_img_tmp['image_raw']], create_graph=True, only_inputs=True) ref_r1_grads_image = ref_r1_grads[0] ref_r1_grads_image_raw = ref_r1_grads[1] ref_r1_penalty_raw = ref_r1_grads_image_raw.square().sum([1,2,3]).mean() ref_r1_penalty_image = ref_r1_grads_image.square().sum([1,2,3]).mean() losses['D_th1kh_gradient_penalty_gt_ref'] = (ref_r1_penalty_image + ref_r1_penalty_raw) / 2 if hparams.get("add_ffhq_singe_disc", False): with conv2d_gradfix.no_weight_gradients(): ref_r1_grads = torch.autograd.grad(outputs=[ffhq_ref_logits.sum()], inputs=[ffhq_ref_img_tmp['image'], ffhq_ref_img_tmp['image_raw']], create_graph=True, only_inputs=True) ref_r1_grads_image = ref_r1_grads[0] ref_r1_grads_image_raw = ref_r1_grads[1] ref_r1_penalty_raw = ref_r1_grads_image_raw.square().sum([1,2,3]).mean() ref_r1_penalty_image = ref_r1_grads_image.square().sum([1,2,3]).mean() losses['D_ffhq_gradient_penalty_gt_ref'] = (ref_r1_penalty_image + ref_r1_penalty_raw) / 2 # pred mv img & D minimize if 'th1kh_recon_mv_imgs' in self.gen_tmp_output: camera = batch['th1kh_mv_cameras'] disc_inp_img = { 'image': self.gen_tmp_output['th1kh_recon_mv_imgs'], 'image_raw': self.gen_tmp_output['th1kh_recon_mv_imgs_raw'], } gen_logits = self.forward_D(disc_inp_img, camera, update_emas=True) losses['D_th1kh_minimize_model_pred_mv'] = torch.nn.functional.softplus(gen_logits).mean() if hparams.get("add_ffhq_singe_disc", False) and (self.global_step+1) % FFHQ_DISC_UPDATE_INTERVAL == 0: gen_logits = self.forward_ffhq_D(disc_inp_img, camera, update_emas=True) losses['D_ffhq_minimize_model_pred_mv'] = torch.nn.functional.softplus(gen_logits).mean() # gt mv img & D maximize mv_cameras = batch['th1kh_mv_cameras'] mv_img_tmp_image = batch['th1kh_mv_head_imgs'].detach().requires_grad_(True) mv_img_tmp_image_raw = batch['th1kh_mv_head_imgs_raw'].detach().requires_grad_(True) th1kh_mv_img_tmp = {'image': mv_img_tmp_image, 'image_raw': mv_img_tmp_image_raw} th1kh_mv_logits = self.forward_D(th1kh_mv_img_tmp, mv_cameras) losses['D_th1kh_maximize_gt_mv'] = torch.nn.functional.softplus(-th1kh_mv_logits).mean() # gt mv img & gradient penalty if (self.global_step + 1) % hparams['reg_interval_d'] == 0 and self.training is True: with conv2d_gradfix.no_weight_gradients(): mv_r1_grads = torch.autograd.grad(outputs=[th1kh_mv_logits.sum()], inputs=[th1kh_mv_img_tmp['image'], th1kh_mv_img_tmp['image_raw']], create_graph=True, only_inputs=True) mv_r1_grads_image = mv_r1_grads[0] mv_r1_grads_image_raw = mv_r1_grads[1] mv_r1_penalty_raw = mv_r1_grads_image_raw.square().sum([1,2,3]).mean() mv_r1_penalty_image = mv_r1_grads_image.square().sum([1,2,3]).mean() losses['D_th1kh_gradient_penalty_gt_mv'] = (mv_r1_penalty_image + mv_r1_penalty_raw) / 2 self.gen_tmp_output = {} return losses def _training_step(self, sample, batch_idx, optimizer_idx): if len(sample) == 0: return None if optimizer_idx == 0: sample = self.prepare_batch(sample) self.cache_sample = sample else: sample = self.cache_sample losses = {} if optimizer_idx == 0: # Train Generator if hparams['two_stage_training']: if self.global_step >= self.start_adv_iters: # only the resolution module requires grad self.model.on_train_superresolution() if hparams.get('also_update_decoder'): self.model.decoder.requires_grad_(True) else: # only the nerf module requires grad self.model.on_train_full_model() else: self.model.on_train_full_model() losses.update(self.run_G_th1kh_src2src_image(sample)) # 提升identity similarity, 很必要, 否则会相似度变差 losses.update(self.run_G_th1kh_src2tgt_image(sample)) losses.update(self.run_G_reg(sample)) losses.update(self.run_G_reg_cond(sample)) loss_weights = { 'G_th1kh_ref_img_mae': hparams.get("lambda_mse", 1.0), 'G_th1kh_ref_img_mae_raw': hparams.get("lambda_mse", 1.0), 'G_th1kh_ref_img_lpips': 0.1, 'G_th1kh_ref_img_lpips_raw': 0.1, 'G_th1kh_ref_adv': hparams['lambda_th1kh_mv_adv'] if self.global_step >= self.start_adv_iters else 0., 'G_th1kh_ref_weights_l1_loss': hparams.get("lambda_weights_l1", 0.5), 'G_th1kh_ref_weights_entropy_loss': hparams.get("lambda_weights_entropy", 0.05), 'G_th1kh_mv_img_mae': hparams.get("lambda_mse", 1.0), 'G_th1kh_mv_img_mae_raw': hparams.get("lambda_mse", 1.0), 'G_th1kh_mv_img_lpips': 0.1, 'G_th1kh_mv_img_lpips_raw': 0.1, 'G_th1kh_mv_adv': hparams['lambda_th1kh_mv_adv'] if self.global_step >= self.start_adv_iters else 0., 'G_th1kh_mv_weights_l1_loss': hparams.get("lambda_weights_l1", 0.3), 'G_th1kh_mv_weights_entropy_loss': hparams.get("lambda_weights_entropy", 0.01), 'G_th1kh_ref_img_lip_mae': 0.5, 'G_th1kh_ref_img_lip_lpips': 0.05, 'G_th1kh_mv_img_lip_mae': 0.5, 'G_th1kh_mv_img_lip_lpips': 0.05, 'G_th1kh_regularize_density_l1': hparams['lambda_density_reg'] * hparams['reg_interval_g'], 'G_ffhq_adv_maxmimize_model_pred_mv': hparams['lambda_ffhq_mv_adv'] if self.global_step >= self.start_adv_iters else 0., 'secc_deform_l1_losses': 0.1, } if 'G_th1kh_regularize_pertube_secc_mae' in losses: target_pertube_blink_secc_loss = hparams.get('target_pertube_blink_secc_loss', 0.15) target_pertube_secc_loss = hparams.get('target_pertube_secc_loss', 0.15) current_pertube_blink_secc_loss = losses['G_th1kh_regularize_blink_secc_mae'].item() current_pertube_secc_loss = losses['G_th1kh_regularize_pertube_secc_mae'].item() grad_lambda_pertube_blink_secc = (math.log10(current_pertube_blink_secc_loss+1e-15) - math.log10(target_pertube_blink_secc_loss+1e-15)) # 如果需要增大lambda_pertube_secc, 则current_loss大于targt, grad值大于0 grad_lambda_pertube_secc = (math.log10(current_pertube_secc_loss+1e-15) - math.log10(target_pertube_secc_loss+1e-15)) # 如果需要增大lambda_pertube_secc, 则current_loss大于targt, grad值大于0 lr_lambda_pertube_secc = hparams.get('lr_lambda_pertube_secc', 0.01) self.model.lambda_pertube_blink_secc.data = self.model.lambda_pertube_blink_secc.data + grad_lambda_pertube_blink_secc * lr_lambda_pertube_secc self.model.lambda_pertube_blink_secc.data.clamp_(0, 2.) self.model.lambda_pertube_secc.data = self.model.lambda_pertube_secc.data + grad_lambda_pertube_secc * lr_lambda_pertube_secc self.model.lambda_pertube_secc.data.clamp_(0, 0.2) if hparams['target_pertube_secc_loss'] == 0.: self.model.lambda_pertube_secc.data = self.model.lambda_pertube_secc.data * 0. if hparams['target_pertube_blink_secc_loss'] == 0.: self.model.lambda_pertube_blink_secc.data = self.model.lambda_pertube_blink_secc.data * 0. losses['lambda_pertube_blink_secc'] = self.model.lambda_pertube_blink_secc.item() losses['lambda_pertube_secc'] = self.model.lambda_pertube_secc.item() loss_weights['G_th1kh_regularize_pertube_secc_mae'] = self.model.lambda_pertube_secc.item() * hparams['reg_interval_g_cond'] # 把新的lambda更新到loss_weights里面 loss_weights['G_th1kh_regularize_blink_secc_mae'] = self.model.lambda_pertube_blink_secc.item() * hparams['reg_interval_g_cond'] # 把新的lambda更新到loss_weights里面 if hparams.get("disable_highreso_at_stage1", False) and hparams['two_stage_training'] and self.global_step >= self.start_adv_iters: loss_weights['G_th1kh_mv_img_mae'] = 0. loss_weights['G_th1kh_mv_img_lpips'] = 0. elif optimizer_idx == 1: # Train Disc if self.global_step < hparams["start_adv_iters"] - 10000: # start train disc too early is a waste of resource return None losses.update(self.forward_D_main(sample)) loss_weights = { 'D_th1kh_maximize_gt_ref': 1.0, 'D_ffhq_maximize_gt_ref': 1.0, 'D_th1kh_maximize_gt_mv': 1.0, 'D_th1kh_minimize_model_pred_mv': 1.0, 'D_ffhq_minimize_model_pred_mv': 1.0, 'D_th1kh_gradient_penalty_gt_ref': hparams['lambda_gradient_penalty'] * hparams['reg_interval_d'], 'D_th1kh_gradient_penalty_gt_mv': hparams['lambda_gradient_penalty'] * hparams['reg_interval_d'], 'D_ffhq_gradient_penalty_gt_ref': hparams['lambda_gradient_penalty'] * hparams['reg_interval_d'], } self.gen_tmp_output = {} else: return None total_loss = sum([loss_weights[k] * v for k, v in losses.items() if isinstance(v, torch.Tensor) and v.requires_grad]) # total_loss = sum([loss_weights.get(k, 1.0) * v for k, v in losses.items() if isinstance(v, torch.Tensor) and v.requires_grad]) if len(losses) == 0: return None return total_loss, losses ##################### # Validation ##################### def validation_start(self): self.gen_dir = os.path.join(hparams['work_dir'], f'validation_results') os.makedirs(self.gen_dir, exist_ok=True) @torch.no_grad() def validation_step(self, sample, batch_idx): self.gen_dir = os.path.join(hparams['work_dir'], f'validation_results') os.makedirs(self.gen_dir, exist_ok=True) outputs = {} losses = {} if len(sample) == 0: return None sample = self.prepare_batch(sample) rank = 0 if len(set(os.environ['CUDA_VISIBLE_DEVICES'].split(","))) == 1 else dist.get_rank() losses.update(self.run_G_th1kh_src2tgt_image(sample)) losses.update(self.run_G_reg(sample)) losses.update(self.run_G_reg_cond(sample)) losses.update(self.forward_D_main(sample)) outputs['losses'] = losses outputs['total_loss'] = sum(outputs['losses'].values()) outputs = tensors_to_scalars(outputs) if self.global_step % hparams['valid_infer_interval'] == 0 \ and batch_idx < hparams['num_valid_plots'] and rank == 0: imgs_ref = sample['th1kh_ref_head_imgs'] gen_img = self.model.forward(imgs_ref, sample['th1kh_mv_cameras'], cond={'cond_cano': sample['th1kh_cano_secc'], 'cond_src': sample['th1kh_ref_secc'], 'cond_tgt': sample['th1kh_mv_secc'], 'ref_head_img': imgs_ref, 'ref_cameras': sample['th1kh_ref_cameras'], 'ref_alphas': sample['th1kh_ref_head_masks'].float(), }, noise_mode='const') gen_img_recon = self.model.forward(imgs_ref, sample['th1kh_ref_cameras'], cond={'cond_cano': sample['th1kh_cano_secc'], 'cond_src': sample['th1kh_ref_secc'], 'cond_tgt': sample['th1kh_ref_secc'], 'ref_head_img': imgs_ref, 'ref_cameras': sample['th1kh_ref_cameras'], 'ref_alphas': sample['th1kh_ref_head_masks'].float(), }, noise_mode='const') imgs_recon = gen_img_recon['image'].permute(0, 2,3,1) imgs_recon_raw = filtered_resizing(gen_img_recon['image_raw'], size=512, f=self.resample_filter, filter_mode='antialiased').permute(0, 2,3,1) imgs_recon_depth = gen_img_recon['image_depth'].permute(0, 2,3,1) imgs_pred_raw = filtered_resizing(gen_img['image_raw'], size=512, f=self.resample_filter, filter_mode='antialiased').permute(0, 2,3,1) imgs_pred = gen_img['image'].permute(0, 2,3,1) imgs_pred_depth = gen_img['image_depth'].permute(0, 2,3,1) imgs_ref = imgs_ref.permute(0,2,3,1) imgs_mv = sample['th1kh_mv_head_imgs'].permute(0,2,3,1) # [B, H, W, 3] for i in range(len(imgs_pred)): idx_string = format(i+batch_idx * hparams['batch_size'], "05d") base_fn = f"{idx_string}" img_ref_mv_recon_pred = torch.cat([imgs_ref[i], imgs_mv[i], imgs_recon_raw[i], imgs_pred_raw[i], imgs_recon[i], imgs_pred[i]], dim=1) ref_secc = filtered_resizing(sample['th1kh_ref_secc'][i].unsqueeze(0), size=512, f=self.resample_filter, filter_mode='antialiased')[0].permute(1,2,0) mv_secc = filtered_resizing(sample['th1kh_mv_secc'][i].unsqueeze(0), size=512, f=self.resample_filter, filter_mode='antialiased')[0].permute(1,2,0) img_ref_mv_recon_pred = torch.cat([img_ref_mv_recon_pred, ref_secc], dim=1) img_ref_mv_recon_pred = torch.cat([img_ref_mv_recon_pred, mv_secc], dim=1) self.save_rgb_to_fname(img_ref_mv_recon_pred, f"{self.gen_dir}/th1kh_images_rgb_iter{self.global_step}/ref_mv_reconraw_predraw_recon_pred_{base_fn}.png") img_depth_recon_pred = torch.cat([imgs_recon_depth[i], imgs_pred_depth[i]], dim=1) self.save_depth_to_fname(img_depth_recon_pred, f"{self.gen_dir}/th1kh_images_depth_iter{self.global_step}/recon_pred_{base_fn}.png") if batch_idx == 0 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0): image_name = "data/raw/examples/Macron.png" imgs_ref = cv2.imread(image_name) img = load_img_to_512_hwc_array(image_name) segmap = self.seg_model._cal_seg_map(img) head_img = self.seg_model._seg_out_img_with_segmap(img, segmap, mode='head')[0] head_mask = segmap[[1,3,5] , :, :].sum(axis=0)[None,:] > 0.5 # glasses 也属于others head_mask = torch.tensor(head_mask).float().reshape([1,1,512,512]).cuda() imgs_ref = ((torch.tensor(head_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w] from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image coeff_dict = fit_3dmm_for_a_image(image_name, save=False) id = torch.tensor(coeff_dict['id']).float().cuda().reshape([1, 80]) exp = torch.tensor(coeff_dict['exp']).float().cuda().reshape([1, 64]) with torch.no_grad(): _, cano_secc = self.secc_renderer(id,exp*0,sample['th1kh_ref_eulers']*0,sample['th1kh_ref_trans']*0) _, ref_secc = self.secc_renderer(id,exp,sample['th1kh_ref_eulers']*0,sample['th1kh_ref_trans']*0) gen_img = self.model.forward(imgs_ref, sample['th1kh_mv_cameras'][0:1], cond={'cond_cano': cano_secc, 'cond_src': ref_secc, 'cond_tgt': ref_secc, 'ref_head_img': imgs_ref, 'ref_cameras': sample['th1kh_mv_cameras'][0:1], 'ref_alphas': head_mask}, noise_mode='const') img = gen_img['image'].permute(0, 2,3,1)[0] self.save_rgb_to_fname(img, f"{self.gen_dir}/ood_images_rgb_iter{self.global_step}/May.png") return outputs def masked_error_loss(self, img_pred, img_gt, mask, unmasked_weight=0.1, mode='l1'): # mask: [B, 1, H, W] # 对raw图像, 因为deform的原因背景没法全黑, 导致这部分mse过高, 我们将其mask掉, 只计算人脸部分 masked_weight = 1.0 weight_mask = mask.float() * masked_weight + (~mask).float() * unmasked_weight if mode == 'l1': error = (img_pred - img_gt).abs().sum(dim=1, keepdim=True) * weight_mask else: error = (img_pred - img_gt).pow(2).sum(dim=1, keepdim=True) * weight_mask error.clamp_(0, max(0.5, error.quantile(0.8).item())) # clamp掉较高loss的pixel, 避免姿态没对齐的pixel导致的异常值占主导影响训练 loss = error.mean() return loss def set_unmasked_to_black(self, img, mask): out_img = img * mask.float() - (~mask).float() # -1 denotes black return out_img def dilate(self, bin_img, ksize=5, mode='max_pool'): """ mode: max_pool or avg_pool """ # bin_img, [b, 1, h, w] pad = (ksize-1)//2 bin_img = F.pad(bin_img, pad=[pad,pad,pad,pad], mode='reflect') if mode == 'max_pool': out = F.max_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0) else: out = F.avg_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0) return out def dilate_mask(self, mask, ksize=21): mask = self.dilate(mask, ksize=ksize, mode='max_pool') return mask def validation_end(self, outputs): return super().validation_end(outputs)