# Main training loop of Portrait4D, modified from EG3D: https://github.com/NVlabs/eg3d # SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual # property and proprietary rights in and to this material, related # documentation and any modifications thereto. Any use, reproduction, # disclosure or distribution of this material and related documentation # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. """Main training loop.""" import os import time import copy import json import pickle import psutil import PIL.Image import numpy as np import random import torch import torch.nn as nn import dnnlib from torch_utils import misc from torch_utils import training_stats from torch_utils.ops import conv2d_gradfix from torch_utils.ops import grid_sample_gradfix from omegaconf import OmegaConf from rendering import RenderingClass from einops import rearrange import legacy from recon.models.lpips.lpips import LPIPS from recon.models.id.id_loss import IDLoss from recon.training.generator.triplane_v20_original import TriPlaneGenerator from headGallery_model.models.triplane_vae import AutoencoderKL as AutoencoderKLTriplane # from training.triplane import PartTriPlaneGeneratorDeform # ---------------------------------------------------------------------------- # def setup_snapshot_image_grid(all_shape_params, all_exp_params, all_pose_params, all_eye_pose_params, all_c, # static_dataset=False, random_seed=1): # gw = 7 # gh = 4 # # grid_indices = np.random.RandomState(random_seed).randint(0, len(all_shape_params), size=(gw * gh)) # # shape_params = all_shape_params[grid_indices] # shape_params = np.tile(np.expand_dims(shape_params, 1), (1, 3, 1)).reshape(gw * gh, 3, -1) # # grid_indices2 = np.random.RandomState(random_seed + 1).randint(0, len(all_exp_params), size=(gw * gh)) # mot_indices = np.random.RandomState(random_seed + 2).randint(0, len(all_exp_params[0]), size=(gw * gh, 2)) # # exp_params = all_exp_params[grid_indices2] # exp_params = np.stack([exp_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # pose_params = all_pose_params[grid_indices2] # pose_params = np.stack([pose_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # eye_pose_params = all_eye_pose_params[grid_indices2] # eye_pose_params = np.stack( # [eye_pose_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # if not static_dataset: # # for dynamic # exp_params = np.concatenate([exp_params, exp_params[:, -1:]], axis=1).reshape(gw * gh, 3, # -1) # (gw * gh, 3, dim) # pose_params = np.concatenate([pose_params, pose_params[:, -1:]], axis=1).reshape(gw * gh, 3, -1) # eye_pose_params = np.concatenate([eye_pose_params, eye_pose_params[:, -1:]], axis=1).reshape(gw * gh, 3, -1) # else: # # for static # exp_params = np.concatenate([exp_params[:, :1], exp_params[:, :1], exp_params[:, :1]], axis=1).reshape(gw * gh, # 3, # -1) # (gw * gh, 3, dim) # pose_params = np.concatenate([pose_params[:, :1], pose_params[:, :1], pose_params[:, :1]], axis=1).reshape( # gw * gh, 3, -1) # eye_pose_params = np.concatenate([eye_pose_params[:, :1], eye_pose_params[:, :1], eye_pose_params[:, :1]], # axis=1).reshape(gw * gh, 3, -1) # # grid_indices3 = np.random.randint(0, len(all_c), size=(gw * gh * 3)) # c = all_c[grid_indices3].reshape(gw * gh, 3, -1) # # return (gw, gh), shape_params, exp_params, pose_params, eye_pose_params, c # def setup_snapshot_image_grid(all_shape_params, all_exp_params, all_pose_params, all_eye_pose_params, all_c, # static_dataset=False, random_seed=1): # gw = 7 # gh = 4 # # grid_indices = np.random.RandomState(random_seed).randint(0, len(all_shape_params), size=(gw * gh)) # # shape_params = all_shape_params[grid_indices] # shape_params = np.tile(np.expand_dims(shape_params, 1), (1, 3, 1)).reshape(gw * gh, 3, -1) # # grid_indices2 = np.random.RandomState(random_seed + 1).randint(0, len(all_exp_params), size=(gw * gh)) # mot_indices = np.random.RandomState(random_seed + 2).randint(0, len(all_exp_params[0]), size=(gw * gh, 2)) # # exp_params = all_exp_params[grid_indices2] # exp_params = np.stack([exp_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # pose_params = all_pose_params[grid_indices2] # pose_params = np.stack([pose_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # eye_pose_params = all_eye_pose_params[grid_indices2] # eye_pose_params = np.stack( # [eye_pose_params[i, mot_indices[i]] for i in range(len(mot_indices))]) # (gw * gh, 2, dim) # # if not static_dataset: # # for dynamic # exp_params = np.concatenate([exp_params, exp_params[:, -1:]], axis=1).reshape(gw * gh, 3, # -1) # (gw * gh, 3, dim) # pose_params = np.concatenate([pose_params, pose_params[:, -1:]], axis=1).reshape(gw * gh, 3, -1) # eye_pose_params = np.concatenate([eye_pose_params, eye_pose_params[:, -1:]], axis=1).reshape(gw * gh, 3, -1) # else: # # for static # exp_params = np.concatenate([exp_params[:, :1], exp_params[:, :1], exp_params[:, :1]], axis=1).reshape(gw * gh, # 3, # -1) # (gw * gh, 3, dim) # pose_params = np.concatenate([pose_params[:, :1], pose_params[:, :1], pose_params[:, :1]], axis=1).reshape( # gw * gh, 3, -1) # eye_pose_params = np.concatenate([eye_pose_params[:, :1], eye_pose_params[:, :1], eye_pose_params[:, :1]], # axis=1).reshape(gw * gh, 3, -1) # # grid_indices3 = np.random.randint(0, len(all_c), size=(gw * gh * 3)) # c = all_c[grid_indices3].reshape(gw * gh, 3, -1) # # return (gw, gh), shape_params, exp_params, pose_params, eye_pose_params, c @torch.no_grad() def setup_snapshot_image_grid_gallery(val_set, vae_triplane, vae_std, vae_mean, render, device): gw = 2 gh = 2 phase_real_z_val, phase_real_latent_val, phase_real_c_1_d_val, phase_real_c_2_d_val, phase_real_c_3_d_val, phase_real_v_1_d_val, phase_real_v_2_d_val, phase_real_v_s_val, motion_1_val, motion_2_val, motion_ffhq_val, model_list_val = next( val_set) phase_real_z_val, phase_real_latent_val, phase_real_c_1_d_val, phase_real_c_2_d_val, phase_real_c_3_d_val, phase_real_v_1_d_val, phase_real_v_2_d_val, phase_real_v_s_val, motion_1_val, motion_2_val, motion_ffhq_val, model_list_val = \ phase_real_z_val.to(device), phase_real_latent_val.to(device), phase_real_c_1_d_val.to(device), phase_real_c_2_d_val.to(device), phase_real_c_3_d_val.to(device), phase_real_v_1_d_val.to(device), phase_real_v_2_d_val.to(device), \ phase_real_v_s_val.to(device), motion_1_val.to(device), motion_2_val.to(device), motion_ffhq_val.to(device), model_list_val batchsize = phase_real_z_val.shape[0] cur_z = phase_real_z_val.unsqueeze(1).repeat(1, 3, 1) # [b,1,512] # cur_latent = phase_real_latent.unsqueeze(1).repeat(1, 1, 1, 1, 1) # [b,3,c, h, w] c_1 = phase_real_c_1_d_val.unsqueeze(1) # input c_2 = phase_real_c_2_d_val.unsqueeze(1) # motion c_3 = phase_real_c_3_d_val.unsqueeze(1) # target cur_c = torch.cat([c_1, c_2, c_3], 1) # from ffhq exp_d_1 = phase_real_v_1_d_val.unsqueeze(1) # motion_1 from vfhq exp_d_2 = phase_real_v_2_d_val.unsqueeze(1) # motion_2 from vfhq exp_s = phase_real_v_s_val.unsqueeze(1).repeat(1, 3, 1, 1, 1) # motion_3 from ffhq for static similar to Portrait4D cur_exp_params = torch.cat([exp_d_1, exp_d_2, exp_d_2], dim=1) cur_exp_out = torch.cat([cur_exp_params[:batchsize // 2], exp_s[batchsize // 2:]], dim=0) model_list_out = [val for val in model_list_val for i in range(3)] model_list_out = listfunc(model_list_out, 3) assert phase_real_latent_val.shape[0] == batchsize cano_tri = vae_triplane.decode(phase_real_latent_val.to(torch.float16)) cano_tri = cano_tri.float() cano_tri = rearrange(cano_tri, "b c f h w -> b f c h w") ref_tri = cano_tri * vae_std + vae_mean exp_target = cur_exp_out[:, 2] # ref_tri_out = render.gen_triplane(ref_tri, exp_target) motion_1_out = torch.cat([motion_1_val[:batchsize // 2], motion_ffhq_val[batchsize // 2:]], dim=0) motion_2_out = torch.cat([motion_2_val[:batchsize // 2], motion_ffhq_val[batchsize // 2:]], dim=0) # ref_img = render.get_img_with_tri(ref_tri, c_3) # always half static return (gw, gh), cur_z, cur_c, cur_exp_out, motion_1_out, motion_2_out, model_list_out, ref_tri # ---------------------------------------------------------------------------- def save_image_grid_all(img_app, img_mot, img_recon, img_ref, fname, drange, grid_size): lo, hi = drange img_app = np.asarray(img_app, dtype=np.float32) img_app = (img_app - lo) * (255 / (hi - lo)) img_app = np.rint(img_app).clip(0, 255).astype(np.uint8) img_mot = np.asarray(img_mot, dtype=np.float32) img_mot = (img_mot - lo) * (255 / (hi - lo)) img_mot = np.rint(img_mot).clip(0, 255).astype(np.uint8) img_recon = np.asarray(img_recon, dtype=np.float32) img_recon = (img_recon - lo) * (255 / (hi - lo)) img_recon = np.rint(img_recon).clip(0, 255).astype(np.uint8) img_ref = np.asarray(img_ref, dtype=np.float32) img_ref = (img_ref - lo) * (255 / (hi - lo)) img_ref = np.rint(img_ref).clip(0, 255).astype(np.uint8) gw, gh = grid_size _N, C, H, W = img_app.shape img = np.concatenate([img_app, img_mot, img_recon, img_ref], -1) gw, gh = grid_size _N, C, H, W = img.shape img = img.reshape([gh, gw, C, H, W]) img = img.transpose(0, 3, 1, 4, 2) img = img.reshape([gh * H, gw * W, C]) assert C in [1, 3] if C == 1: PIL.Image.fromarray(img[:, :, 0], 'L').save(fname) if C == 3: PIL.Image.fromarray(img, 'RGB').save(fname) def save_image_grid(img, fname, drange, grid_size): lo, hi = drange img = np.asarray(img, dtype=np.float32) img = (img - lo) * (255 / (hi - lo)) img = np.rint(img).clip(0, 255).astype(np.uint8) gw, gh = grid_size _N, C, H, W = img.shape img = img.reshape([gh, gw, C, H, W]) img = img.transpose(0, 3, 1, 4, 2) img = img.reshape([gh * H, gw * W, C]) assert C in [1, 3] if C == 1: PIL.Image.fromarray(img[:, :, 0], 'L').save(fname) if C == 3: PIL.Image.fromarray(img, 'RGB').save(fname) def set_requires_grad(nets, requires_grad=False): if not isinstance(nets, list): nets = [nets] for net in nets: if net is not None: for param in net.parameters(): param.requires_grad = requires_grad def split_gen(gen, batch_gpu, batch_size, device): assert type(gen) == list if type(gen[0]) == np.ndarray: all_gen = torch.from_numpy(np.stack(gen)).pin_memory().to(device).float() all_gen = [phase_gen_c.split(batch_gpu) for phase_gen_c in all_gen.split(batch_size)] elif type(gen[0]) == dict: all_gen = [[{} for _ in range(batch_size // batch_gpu)] for _ in range(len(gen) // batch_size)] for key in gen[0].keys(): key_value = torch.from_numpy(np.stack([sub[key] for sub in gen])).pin_memory().to(device).float() key_value_split = [phase_gen_c.split(batch_gpu) for phase_gen_c in key_value.split(batch_size)] for i in range(len(key_value_split)): for j in range(len(key_value_split[i])): all_gen[i][j][key] = key_value_split[i][j] else: raise NotImplementedError return all_gen def split_gen_new(gen, batch_gpu, batch_size, device): if type(gen) == torch.Tensor: all_gen = gen.view((gen.shape[0] * gen.shape[1],) + (gen.shape[2:])).pin_memory().to(device) all_gen = [phase_gen_c.split(batch_gpu) for phase_gen_c in all_gen.split(batch_size)] elif type(gen[0]) == dict: all_gen = [[{} for _ in range(batch_size // batch_gpu)] for _ in range(int(len(gen) * list(gen[0].values())[0].shape[0] // batch_size))] for key in gen[0].keys(): key_value = torch.cat([sub[key] for sub in gen], dim=0).pin_memory().to(device) key_value_split = [phase_gen_c.split(batch_gpu) for phase_gen_c in key_value.split(batch_size)] for i in range(len(key_value_split)): for j in range(len(key_value_split[i])): all_gen[i][j][key] = key_value_split[i][j] else: raise NotImplementedError return all_gen @torch.no_grad() # similar to the fetch_random_params def fetch_dataset(phase_real_z, phase_real_latent, phase_real_c_1_d, phase_real_c_2_d, phase_real_c_3_d, phase_real_v_1_d, phase_real_v_2_d, phase_real_v_s, motion_1, motion_2, motion_ffhq, model_list, vae_triplane, vae_std, vae_mean, render): batchsize = phase_real_z.shape[0] cur_z = phase_real_z.unsqueeze(1).repeat(1, 3, 1) # [b,1,512] # cur_latent = phase_real_latent.unsqueeze(1).repeat(1, 1, 1, 1, 1) # [b,3,c, h, w] c_1 = phase_real_c_1_d.unsqueeze(1) c_2 = phase_real_c_2_d.unsqueeze(1) c_3 = phase_real_c_3_d.unsqueeze(1) cur_c = torch.cat([c_1, c_2, c_3], 1) # from ffhq exp_d_1 = phase_real_v_1_d.unsqueeze(1) # motion_1 from vfhq exp_d_2 = phase_real_v_2_d.unsqueeze(1) # motion_2 from vfhq exp_s = phase_real_v_s.unsqueeze(1).repeat(1, 3, 1, 1, 1) # motion_3 from ffhq for static similar to Portrait4D cur_exp_params = torch.cat([exp_d_1, exp_d_2, exp_d_2], dim=1) cur_exp_out = torch.cat([cur_exp_params[:batchsize // 2], exp_s[batchsize // 2:]], dim=0) model_list_out = [val for val in model_list for i in range(3)] model_list_out = listfunc(model_list_out, 3) assert phase_real_latent.shape[0] == batchsize cano_tri = vae_triplane.decode(phase_real_latent.to(torch.float16)) cano_tri = cano_tri.float() cano_tri = rearrange(cano_tri, "b c f h w -> b f c h w") ref_tri = cano_tri * vae_std + vae_mean exp_target = cur_exp_out[:, 2] # ref_tri_out = render.gen_triplane(ref_tri, exp_target) motion_1_out = torch.cat([motion_1[:batchsize // 2], motion_ffhq[batchsize // 2:]], dim=0) motion_2_out = torch.cat([motion_2[:batchsize // 2], motion_ffhq[batchsize // 2:]], dim=0) # always half static return cur_z, cur_c, cur_exp_out, motion_1_out, motion_2_out, model_list_out, cano_tri, ref_tri, exp_target # choose random FLAME parameters for online data synthesis (torch version) def gan_model(gan_models, device, gan_model_base_dir): gan_model_dict = gan_models["gan_models"] gan_model_load = {} for model_name in gan_model_dict.keys(): model_pkl = os.path.join(gan_model_base_dir, model_name + '.pkl') with dnnlib.util.open_url(model_pkl) as f: G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore G_new = TriPlaneGenerator(*G.init_args, **G.init_kwargs).eval().requires_grad_(False).to(device) misc.copy_params_and_buffers(G, G_new, require_all=True) G_new.neural_rendering_resolution = G.neural_rendering_resolution G_new.rendering_kwargs = G.rendering_kwargs gan_model_load[model_name] = G_new return gan_model_load def listfunc(listTemp, n): out_list = [] for i in range(0, len(listTemp), n): each = listTemp[i:i + n] out_list.append(each) return out_list def collate_fn(data): model_list = [example["model_name"] for example in data] phase_real_z = torch.cat([example["phase_real_z"] for example in data], dim=0) phase_real_latent = torch.cat([example["phase_real_latent"] for example in data], dim=0) phase_real_c_1_d = torch.cat([example["phase_real_c_1_d"] for example in data], dim=0) phase_real_c_2_d = torch.cat([example["phase_real_c_2_d"] for example in data], dim=0) phase_real_c_3_d = torch.cat([example["phase_real_c_3_d"] for example in data], dim=0) phase_real_v_s = torch.cat([example["phase_real_v_s"] for example in data], dim=0) motion_ffhq = torch.cat([example["motion_ffhq"] for example in data], dim=0) motion_1 = torch.cat([example["motion_1"] for example in data], dim=0) motion_2 = torch.cat([example["motion_2"] for example in data], dim=0) phase_real_v_1_d = torch.cat([example["phase_real_v_1_d"] for example in data], dim=0) phase_real_v_2_d = torch.cat([example["phase_real_v_2_d"] for example in data], dim=0) return phase_real_z, phase_real_latent, phase_real_c_1_d, phase_real_c_2_d, phase_real_c_3_d, phase_real_v_1_d, phase_real_v_2_d, phase_real_v_s, motion_1, motion_2, motion_ffhq, model_list # ---------------------------------------------------------------------------- def training_loop( run_dir='.', # Output directory. training_set_kwargs={}, # Options for training set. data_loader_kwargs={}, # Options for torch.utils.data.DataLoader. G_kwargs={}, # Options for generator network. D_kwargs={}, # Options for discriminator network. D_patch_kwargs={}, # Options for patch discriminator (deprecated). G_opt_kwargs={}, # Options for generator optimizer. D_opt_kwargs={}, # Options for discriminator optimizer. D_patch_opt_kwargs={}, # Options for patch discriminator optimizer (deprecated). augment_kwargs=None, # Options for augmentation pipeline. None = disable. loss_kwargs={}, # Options for loss function. metrics=[], # Metrics to evaluate during training. random_seed=0, # Global random seed. num_gpus=1, # Number of GPUs participating in the training. rank=0, # Rank of the current process in [0, num_gpus[. batch_size=4, # Total batch size for one training iteration. Can be larger than batch_gpu * num_gpus. batch_gpu=4, # Number of samples processed at a time by one GPU. ema_kimg=10, # Half-life of the exponential moving average (EMA) of generator weights. ema_rampup=0.05, # EMA ramp-up coefficient. None = no rampup. G_reg_interval=None, # How often to perform regularization for G? None = disable lazy regularization. D_reg_interval=16, # How often to perform regularization for D? None = disable lazy regularization. D_patch_reg_interval=16, # How often to perform regularization for D patch (deprecated) augment_p=0, # Initial value of augmentation probability. ada_target=None, # ADA target value. None = fixed p. ada_interval=4, # How often to perform ADA adjustment? ada_kimg=500, # ADA adjustment speed, measured in how many kimg it takes for p to increase/decrease by one unit. total_kimg=25000, # Total length of the training, measured in thousands of real images. kimg_per_tick=4, # Progress snapshot interval. image_snapshot_ticks=50, # How often to save image snapshots? None = disable. network_snapshot_ticks=50, # How often to save network snapshots? None = disable. resume_pkl=None, # Network pickle to resume training from. resume_kimg=0, # First kimg to report when resuming training. cudnn_benchmark=True, # Enable torch.backends.cudnn.benchmark? abort_fn=None, # Callback function for determining whether to abort training. Must return consistent results across ranks. progress_fn=None, # Callback function for updating training progress. Called for all ranks. motion_scale=1.0, # Scale of the motion-related cross-attention outputs. swapping_prob=0.5, # Probability to set dynamic data as static data. half_static=True, # Whether or not to set the second half of the batchsize as static data. resume_pkl_G_syn=None, # Checkpoint of pre-trained GenHead generator for training data synthesis. truncation_psi=0.7, # Truncation rate of GenHead for training data synthesis. cross_lr_scale=1.0, # Learning rate scale of the motion-related cross attentions. gan_model_base_dir=None, vae_pretrained=None, render_pretrain=None, vae_triplane_config=None, pretrain_portrait_4D = None, load_tri_pretrain = True, ): # Initialize. start_time = time.time() device = torch.device('cuda', rank) np.random.seed(random_seed * num_gpus + rank) torch.manual_seed(random_seed * num_gpus + rank) torch.backends.cudnn.benchmark = cudnn_benchmark # Improves training speed. torch.backends.cuda.matmul.allow_tf32 = False # Improves numerical accuracy. torch.backends.cudnn.allow_tf32 = False # Improves numerical accuracy. torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False # Improves numerical accuracy. conv2d_gradfix.enabled = True # Improves training speed. # TODO: ENABLE grid_sample_gradfix.enabled = False # Avoids errors with the augmentation pipe. # Load training set. batch_size_dataset = batch_size // num_gpus if rank == 0: print('Loading training set...') training_set = dnnlib.util.construct_class_by_name( **training_set_kwargs) # subclass of training.dataset.Dataset # Construct networks. if rank == 0: print('Constructing networks...') common_kwargs = dict(c_dim=25, img_resolution=512, img_channels=3) G = dnnlib.util.construct_class_by_name(**G_kwargs, **common_kwargs).train().requires_grad_(False).to( device) # subclass of torch.nn.Module for m in G.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() D_img_channel = training_set.num_channels * 3 if G_kwargs.rendering_kwargs.get( 'gen_lms_cond') else training_set.num_channels * 2 # if G_kwargs.rendering_kwargs.get('gen_mask_cond'): D_img_channel += 1 D = dnnlib.util.construct_class_by_name(c_dim=25, img_resolution=512, img_channels=D_img_channel, **D_kwargs).train().requires_grad_(False).to( device) # subclass of torch.nn.Module # Load pre-trained GenHead model if rank == 0: print(f'Resuming GenHead from "{resume_pkl_G_syn}"') config_gan_model = OmegaConf.load(resume_pkl_G_syn) G_syn_dict = gan_model(config_gan_model, device, gan_model_base_dir) # with dnnlib.util.open_url(resume_pkl_G_syn) as f: # G_syn = legacy.load_network_pkl(f)['G_ema'].eval().requires_grad_(False).to(device) # G_syn = PartTriPlaneGeneratorDeform(*G_syn_meta.init_args, **G_syn_meta.init_kwargs).eval().requires_grad_(False).to(device) # misc.copy_params_and_buffers(G_syn_meta, G_syn, require_all=False) # G_syn.neural_rendering_resolution = G_syn_meta.neural_rendering_resolution # G_syn.rendering_kwargs = G_syn_meta.rendering_kwargs # For VAE decoder config_vae_triplane = OmegaConf.load(vae_triplane_config) vae_triplane = AutoencoderKLTriplane(ddconfig=config_vae_triplane['ddconfig'], lossconfig=None, embed_dim=8) vae_triplane_model_file = os.path.join(vae_pretrained, 'pytorch_model.bin') if not os.path.isfile(vae_triplane_model_file): raise RuntimeError(f"{vae_triplane_model_file} does not exist") vae_triplane_state_dict = torch.load(vae_triplane_model_file, map_location="cpu") vae_triplane.load_state_dict(vae_triplane_state_dict) vae_triplane.requires_grad_(False) vae_triplane = vae_triplane.to(device, dtype=torch.float16) # For rendering Rendering = RenderingClass(device, config_vae_triplane['render_network_pkl'], ws_avg_pkl=config_vae_triplane['ws_avg']) data_std = torch.load(config_vae_triplane['std_dir']).to(device).reshape(1, -1, 1, 1, 1) data_mean = torch.load(config_vae_triplane['mean_dir']).to(device).reshape(1, -1, 1, 1, 1) # For LPIPS loss computation lpips = LPIPS(vgg_path=config_vae_triplane['vgg_path'], net='vgg', model_path=config_vae_triplane['vgg']).to(device) set_requires_grad(lpips, requires_grad=False) # For ID loss computation idloss = IDLoss(config_vae_triplane['ir_se50']) idloss = idloss.eval().to(device) set_requires_grad(idloss, requires_grad=False) # For PD-FGC motion embedding extraction # pd_fgc = FanEncoder() # weight_dict = torch.load(motion_pretrained) # pd_fgc.load_state_dict(weight_dict, strict=False) # pd_fgc = pd_fgc.eval().to(device) # set_requires_grad(pd_fgc, requires_grad=False) # set D_patch for 3D-to-2D imitation (deprecated), see Mimic3D for details: https://github.com/SeanChenxy/Mimic3D check this D_patch = None # if loss_kwargs.patch_scale < 1: # img_resolution = loss_kwargs.neural_rendering_resolution_initial if loss_kwargs.neural_rendering_resolution_final is None else loss_kwargs.neural_rendering_resolution_final # common_patch_kwargs = dict(c_dim=0, img_resolution=img_resolution, img_channels=3) # D_patch = dnnlib.util.construct_class_by_name(**D_patch_kwargs, **common_patch_kwargs).train().requires_grad_( # False).to(device) # subclass of torch.nn.Module # if pretrain_portrait_4D is not None and resume_pkl is None: # print(f'Resuming encoders from "{pretrain_portrait_4D}"') # with dnnlib.util.open_url(pretrain_portrait_4D) as f: # print("Reloading Modules!") # load_model = legacy.load_network_pkl(f) # G_encoder = load_model['G_ema'] # misc.copy_params_and_buffers(G_encoder.encoder_global, G.encoder_global ) # misc.copy_params_and_buffers(G_encoder.encoder_detail, G.encoder_detail ) # del G_encoder if (resume_pkl is None) and (rank == 0): print(f'Resuming rendering and super and D from "{render_pretrain}"') with dnnlib.util.open_url(render_pretrain) as f: print("Reloading Modules!") load_model = legacy.load_network_pkl(f) G_deco = load_model['G_ema'] D_deco = load_model['D'] misc.copy_params_and_buffers(G_deco.decoder, G.decoder, require_all=True) misc.copy_params_and_buffers(G_deco.renderer, G.renderer, require_all=True) misc.copy_params_and_buffers(G_deco.ray_sampler, G.ray_sampler, require_all=True) misc.copy_params_and_buffers(G_deco.superresolution, G.superresolution, require_all=True) misc.copy_params_and_buffers(D_deco, D, require_all=True) # if load_tri_pretrain: # misc.copy_params_and_buffers(Rendering.Render.face_backbone, G.face_backbone, require_all=True) # misc.copy_params_and_buffers(Rendering.Render.triplnae_encoder, G.triplnae_encoder, require_all=True) G_ema = copy.deepcopy(G).eval() # Resume from existing pickle. if (resume_pkl is not None) and (rank == 0): print(f'Resuming from "{resume_pkl}"') with dnnlib.util.open_url(resume_pkl) as f: resume_data = legacy.load_network_pkl(f) load_model = [('G', G), ('G_ema', G_ema)] if D is not None: load_model.append(('D', D)) if D_patch is not None: load_model.append(('D_patch', D_patch)) for name, module in load_model: if name in resume_data and resume_data[name] is not None: misc.copy_params_and_buffers(resume_data[name], module, require_all=False) else: print(f'resume_data do not have {name}') if 'augment_pipe' in resume_data and resume_data['augment_pipe'] is not None: augment_p = resume_data['augment_pipe'].p # Setup augmentation. if rank == 0: print('Setting up augmentation...') augment_pipe = None ada_stats = None if (augment_kwargs is not None) and (augment_p > 0 or ada_target is not None): augment_pipe = dnnlib.util.construct_class_by_name(**augment_kwargs).train().requires_grad_(False).to( device) # subclass of torch.nn.Module augment_pipe.p.copy_(torch.as_tensor(augment_p)) if ada_target is not None: ada_stats = training_stats.Collector(regex='Loss/signs/real') # Distribute across GPUs. if rank == 0: print(f'Distributing across {num_gpus} GPUs...') for module in [G, D, G_ema, augment_pipe, lpips, D_patch]: if module is not None: for param in misc.params_and_buffers(module): if param.numel() > 0 and num_gpus > 1: torch.distributed.broadcast(param, src=0) # Setup training phases. if rank == 0: print('Setting up training phases...') # conditioning_params = torch.load(config_vae_triplane['conditioning_params_dir']).to(device) loss = dnnlib.util.construct_class_by_name(device=device, G=G, D=D, G_syn=G_syn_dict, D_patch=D_patch, augment_pipe=augment_pipe, lpips=lpips, id_loss=idloss, conditioning_params=config_vae_triplane['conditioning_params_dir'], w_avg=config_vae_triplane['ws_avg'], **loss_kwargs) # subclass of training.loss.Loss phases = [] phases_asserts = [('G', G, G_opt_kwargs, G_reg_interval), ] if D is not None: phases_asserts.append(('D', D, D_opt_kwargs, D_reg_interval)) # no d_patch if D_patch is not None: phases_asserts.append(('D_patch', D_patch, D_patch_opt_kwargs, D_patch_reg_interval)) for name, module, opt_kwargs, reg_interval in phases_asserts: # if G_update_all is False: # parameter_names = [n for (n, p) in module.named_parameters() if 'superresolution' not in n and not ('decoder' in n and 'encoder_global' not in n) and 'bn' not in n] # do not update mlp and super-resolution following Real-Time Radiance Fields for Single-Image Portrait View Synthesis # else: parameter_names = [n for (n, p) in module.named_parameters() if 'bn' not in n] if name == 'G': parameters_group = [] parameters_cross_names = [n for n in parameter_names if 'encoder_canonical' in n and ( 'maps' in n or 'maps_neutral' in n or 'proj_y' in n or 'proj_y_neutral' in n or 'norm2' in n or 'attn2' in n)] parameters_base_names = [n for n in parameter_names if not n in parameters_cross_names] parameters_cross = [p for (n, p) in module.named_parameters() if n in parameters_cross_names] parameters_base = [p for (n, p) in module.named_parameters() if n in parameters_base_names] parameters_group.append({'params': parameters_cross, 'name': 'G_cross'}) parameters_group.append({'params': parameters_base, 'name': 'G_base'}) parameters = parameters_group else: parameters = [p for (n, p) in module.named_parameters() if n in parameter_names] if reg_interval is None: opt = dnnlib.util.construct_class_by_name(parameters, **opt_kwargs) # subclass of torch.optim.Optimizer phases += [dnnlib.EasyDict(name=name + 'both', module=module, opt=opt, interval=1)] else: # Lazy regularization. mb_ratio = reg_interval / (reg_interval + 1) opt_kwargs = dnnlib.EasyDict(opt_kwargs) opt_kwargs.lr = opt_kwargs.lr * mb_ratio opt_kwargs.betas = [beta ** mb_ratio for beta in opt_kwargs.betas] opt = dnnlib.util.construct_class_by_name(parameters, **opt_kwargs) # subclass of torch.optim.Optimizer phases += [dnnlib.EasyDict(name=name + 'main', module=module, opt=opt, interval=1)] phases += [dnnlib.EasyDict(name=name + 'reg', module=module, opt=opt, interval=reg_interval)] if name == 'G': for param_group in opt.param_groups: if param_group['name'] == 'G_cross': param_group['lr'] = param_group['lr'] * cross_lr_scale for phase in phases: phase.start_event = None phase.end_event = None if rank == 0: phase.start_event = torch.cuda.Event(enable_timing=True) phase.end_event = torch.cuda.Event(enable_timing=True) # Export sample images. grid_size = None grid_z = None grid_c = None training_set_sampler = misc.InfiniteSampler(dataset=training_set, rank=rank, num_replicas=num_gpus, seed=random_seed) training_set_iterator = iter( torch.utils.data.DataLoader(dataset=training_set, sampler=training_set_sampler, collate_fn=collate_fn, batch_size=batch_size // num_gpus, **data_loader_kwargs)) val_set_iterator = iter( torch.utils.data.DataLoader(dataset=training_set ,collate_fn=collate_fn, shuffle=True, batch_size=4, num_workers=1)) if rank == 0: print('Exporting sample images...') batch_gpu_val = 4 grid_size, cur_z_vals, cur_c_vals, cur_exp_out_vals, motion_1_out_vals, motion_2_out_vals, model_list_out_vals, ref_tri_vals = setup_snapshot_image_grid_gallery( val_set_iterator, vae_triplane, data_std, data_mean, Rendering, device) grid_size = (2, 2) cur_z_vals = cur_z_vals.split(batch_gpu_val) cur_c_vals = cur_c_vals.split(batch_gpu_val) cur_exp_out_vals = cur_exp_out_vals.split(batch_gpu_val) model_list_out_vals = listfunc(model_list_out_vals, batch_gpu_val) motion_1_out_vals = motion_1_out_vals.split(batch_gpu_val) motion_2_out_vals = motion_2_out_vals.split(batch_gpu_val) ref_tri_vals = ref_tri_vals.split(batch_gpu_val) out = [] ref_out = [] grid_c = [] grid_c_recon = [] kkkk = 0 with torch.no_grad(): for cur_z_val, cur_exp_out_val, cur_c_val, model_list_out_val, ref_tri_val in \ zip(cur_z_vals, cur_exp_out_vals, cur_c_vals, model_list_out_vals, ref_tri_vals): syn_out = loss.gen_data_by_G_next3D(cur_z_val, cur_exp_out_val, cur_c_val, model_list_out_val) ref_tri_tri = Rendering.gen_triplane(ref_tri_val, cur_exp_out_val[:,2]) ref_imgs_out = Rendering.get_img_with_tri(ref_tri_tri, syn_out['c'][:,2]) grid_c_recon.append(syn_out['c'][:,2]) out.append(syn_out) ref_out.append(ref_imgs_out) grid_c_recon.append(torch.load(config_vae_triplane['syn_out_c_path']).to(device)) images_all = torch.cat([o['image_sr'] for o in out], dim=0) ref_imgs_all_val = torch.cat(ref_out, dim=0) images_all = images_all.reshape(-1, 3, *images_all.shape[-3:]) images_app_val = images_all[:, 0] images_mot_val = images_all[:, 1] images_recon_val = images_all[:, 2] save_image_grid_all(127.5 * (images_app_val.cpu().numpy() + 1), 127.5 * (images_mot_val.cpu().numpy() + 1), 127.5 * (images_recon_val.cpu().numpy() + 1), 127.5 * (ref_imgs_all_val.cpu().numpy() + 1), os.path.join(run_dir, 'reals.png'), drange=[0, 255], grid_size=grid_size) # grid_c = torch.cat(grid_c) images_app_val = images_app_val.split(batch_gpu_val) images_mot_val = images_mot_val.split(batch_gpu_val) # grid_c = grid_c.split(batch_gpu) # Initialize logs. if rank == 0: print('Initializing logs...') stats_collector = training_stats.Collector(regex='.*') stats_metrics = dict() stats_jsonl = None stats_tfevents = None if rank == 0: stats_jsonl = open(os.path.join(run_dir, 'stats.jsonl'), 'wt') try: import torch.utils.tensorboard as tensorboard stats_tfevents = tensorboard.SummaryWriter(run_dir) except ImportError as err: print('Skipping tfevents export:', err) # Train. if rank == 0: print(f'Training for {total_kimg} kimg...') print() cur_nimg = resume_kimg * 1000 cur_tick = 0 tick_start_nimg = cur_nimg tick_start_time = time.time() maintenance_time = tick_start_time - start_time batch_idx = 0 if progress_fn is not None: progress_fn(0, total_kimg) batch_num = batch_size // num_gpus while True: # Fetch training data. with torch.autograd.profiler.record_function('data_fetch'): phase_real_z, phase_real_latent, phase_real_c_1_d, phase_real_c_2_d, phase_real_c_3_d, phase_real_v_1_d, phase_real_v_2_d, phase_real_v_s, motion_1, motion_2, motion_ffhq, model_list = next( training_set_iterator) phase_real_z, phase_real_c, phase_real_exp_params, motion_1, motion_2, model_list, phase_real_cano_tri, phase_real_ref_tri, phase_real_exp_params_target = fetch_dataset( phase_real_z.to(device), phase_real_latent.to(device), phase_real_c_1_d.to(device), phase_real_c_2_d.to(device), phase_real_c_3_d.to(device), phase_real_v_1_d.to(device), phase_real_v_2_d.to(device), phase_real_v_s.to(device), motion_1.to(device), motion_2.to(device), motion_ffhq.to(device), model_list, vae_triplane, data_std, data_mean, Rendering) phase_real_z = phase_real_z.split(batch_gpu) # phase_real_latent = phase_real_latent.split(batch_gpu).to(device) phase_real_c = phase_real_c.split(batch_gpu) phase_real_exp_params = phase_real_exp_params.split(batch_gpu) phase_real_motions_app = motion_1.split(batch_gpu) phase_real_motions = motion_2.split(batch_gpu) phase_real_cano_tri = phase_real_cano_tri.split(batch_gpu) phase_real_ref_tri = phase_real_ref_tri.split(batch_gpu) phase_real_model_list = listfunc(model_list, batch_gpu) phase_real_exp_params_target = phase_real_exp_params_target.split(batch_gpu) # --------------------------------------------------------------------------------------------------------------------------------------- # Online data generation. For efficiency, use same generated data for different phases phase_real_img_app = [] phase_real_img_mot = [] phase_real_img_recon = [] phase_real_depth_recon = [] phase_real_feature_recon = [] phase_real_triplane_recon = [] phase_real_c_recon = [] # phase_real_motions_app = [] # phase_real_motions = [] # phase_real_triplane_latent_recon = [] with torch.no_grad(): for real_z, real_pose_params, real_exp_params, real_models in \ zip(phase_real_z, phase_real_c, phase_real_exp_params, phase_real_model_list): syn_out = loss.gen_data_by_G_next3D(real_z, real_exp_params, real_pose_params, real_models) # Multiview images real_img = syn_out['image_sr'] real_img = real_img.reshape(-1, 3, *real_img.shape[1:]) real_img_app = real_img[:, 0] real_img_mot = real_img[:, 1] real_img_recon = real_img[:, 2] # Segmentation masks # Camera poses real_c_recon = syn_out['c'] # real_c_recon = real_c_recon.reshape(-1, 3, *real_c_recon.shape[1:]) real_c_recon = real_c_recon[:, 2] # Depth images real_depth_recon = syn_out['image_depth'] real_depth_recon = real_depth_recon.reshape(-1, 3, *real_depth_recon.shape[1:]) real_depth_recon = real_depth_recon[:, 2] # Feature maps before super-resolution module real_feature_recon = syn_out['image_feature'] real_feature_recon = real_feature_recon.reshape(-1, 3, *real_feature_recon.shape[1:]) real_feature_recon = real_feature_recon[:, 2] # Sampled tri-plane features real_triplane_recon = syn_out['triplane'] real_triplane_recon = real_triplane_recon.reshape(-1, 3, *real_triplane_recon.shape[1:]) real_triplane_recon = real_triplane_recon[:, 2] # Sampled latent recon # real_triplane_latent_recon = syn_out['vae_out_tri'] phase_real_img_app.append(real_img_app) phase_real_img_mot.append(real_img_mot) phase_real_img_recon.append(real_img_recon) phase_real_depth_recon.append(real_depth_recon) phase_real_feature_recon.append(real_feature_recon) phase_real_triplane_recon.append(real_triplane_recon) phase_real_c_recon.append(real_c_recon) # phase_real_motions_app.append(real_motions_app) # phase_real_motions.append(real_motions) # phase_real_triplane_latent_recon.append(real_triplane_latent_recon) # Execute training phases. for phase in phases: if batch_idx % phase.interval != 0: continue if phase.start_event is not None: phase.start_event.record(torch.cuda.current_stream(device)) # Accumulate gradients. phase.opt.zero_grad(set_to_none=True) phase.module.requires_grad_(True) for real_img_app, real_img_mot, real_img_recon, real_depth_recon, real_feature_recon, real_triplane_recon, real_c_recon, real_motions_app, real_motions, real_cano_tri, real_ref_tri, real_exp_target in \ zip(phase_real_img_app, phase_real_img_mot, phase_real_img_recon, phase_real_depth_recon, phase_real_feature_recon, phase_real_triplane_recon, phase_real_c_recon, phase_real_motions_app, phase_real_motions, phase_real_cano_tri, phase_real_ref_tri, phase_real_exp_params_target): loss.accumulate_gradients(phase=phase.name, real_img_app=real_img_app, real_img_mot=real_img_mot, real_img_recon=real_img_recon, real_depth_recon=real_depth_recon, real_feature_recon=real_feature_recon, real_triplane_recon=real_triplane_recon, real_c_recon=real_c_recon, mesh= real_exp_target, motions_app=real_motions_app, motions=real_motions, real_cano_tri=real_cano_tri, real_ref_tri=real_ref_tri, gain=phase.interval, cur_nimg=cur_nimg, motion_scale=motion_scale, swapping_prob=swapping_prob, half_static=half_static) phase.module.requires_grad_(False) # Update weights. with torch.autograd.profiler.record_function(phase.name + '_opt'): # Do not update mlp decoder and super-resolution module at the warm-up stage following Live3dportrait: https://arxiv.org/abs/2305.02310 if cur_nimg <= loss.discrimination_kimg * 1e3 and phase.name == 'G': sub_params = [p for (n, p) in phase.module.named_parameters() if 'superresolution' in n or ('decoder' in n and 'encoder_global' not in n)] for param in sub_params: if param.grad is not None: param.grad.zero_() params = [param for param in phase.module.parameters() if param.numel() > 0 and param.grad is not None] if len(params) > 0: flat = torch.cat([param.grad.flatten() for param in params]) if num_gpus > 1: torch.distributed.all_reduce(flat) flat /= num_gpus misc.nan_to_num(flat, nan=0, posinf=1e5, neginf=-1e5, out=flat) grads = flat.split([param.numel() for param in params]) for param, grad in zip(params, grads): param.grad = grad.reshape(param.shape) phase.opt.step() # Phase done. if phase.end_event is not None: phase.end_event.record(torch.cuda.current_stream(device)) # Update G_ema. with torch.autograd.profiler.record_function('Gema'): ema_nimg = ema_kimg * 1000 if ema_rampup is not None: ema_nimg = min(ema_nimg, cur_nimg * ema_rampup) ema_beta = 0.5 ** (batch_size / max(ema_nimg, 1e-8)) for p_ema, p in zip(G_ema.parameters(), G.parameters()): p_ema.copy_(p.lerp(p_ema, ema_beta)) for b_ema, b in zip(G_ema.buffers(), G.buffers()): b_ema.copy_(b) G_ema.neural_rendering_resolution = G.neural_rendering_resolution G_ema.rendering_kwargs = G.rendering_kwargs.copy() # Update state. cur_nimg += batch_size batch_idx += 1 # Execute ADA heuristic. if (ada_stats is not None) and (batch_idx % ada_interval == 0): ada_stats.update() adjust = np.sign(ada_stats['Loss/signs/real'] - ada_target) * (batch_size * ada_interval) / ( ada_kimg * 1000) augment_pipe.p.copy_((augment_pipe.p + adjust).max(misc.constant(0, device=device))) # Perform maintenance tasks once per tick. done = (cur_nimg >= total_kimg * 1000) if (not done) and (cur_tick != 0) and (cur_nimg < tick_start_nimg + kimg_per_tick * 1000): continue # Print status line, accumulating the same information in training_stats. tick_end_time = time.time() fields = [] fields += [f"tick {training_stats.report0('Progress/tick', cur_tick):<5d}"] fields += [f"kimg {training_stats.report0('Progress/kimg', cur_nimg / 1e3):<8.1f}"] fields += [ f"time {dnnlib.util.format_time(training_stats.report0('Timing/total_sec', tick_end_time - start_time)):<12s}"] fields += [f"sec/tick {training_stats.report0('Timing/sec_per_tick', tick_end_time - tick_start_time):<7.1f}"] fields += [ f"sec/kimg {training_stats.report0('Timing/sec_per_kimg', (tick_end_time - tick_start_time) / (cur_nimg - tick_start_nimg) * 1e3):<7.2f}"] fields += [f"maintenance {training_stats.report0('Timing/maintenance_sec', maintenance_time):<6.1f}"] fields += [ f"cpumem {training_stats.report0('Resources/cpu_mem_gb', psutil.Process(os.getpid()).memory_info().rss / 2 ** 30):<6.2f}"] fields += [ f"gpumem {training_stats.report0('Resources/peak_gpu_mem_gb', torch.cuda.max_memory_allocated(device) / 2 ** 30):<6.2f}"] fields += [ f"reserved {training_stats.report0('Resources/peak_gpu_mem_reserved_gb', torch.cuda.max_memory_reserved(device) / 2 ** 30):<6.2f}"] torch.cuda.reset_peak_memory_stats() fields += [ f"augment {training_stats.report0('Progress/augment', float(augment_pipe.p.cpu()) if augment_pipe is not None else 0):.3f}"] training_stats.report0('Timing/total_hours', (tick_end_time - start_time) / (60 * 60)) training_stats.report0('Timing/total_days', (tick_end_time - start_time) / (24 * 60 * 60)) if rank == 0: print(' '.join(fields)) # Check for abort. if (not done) and (abort_fn is not None) and abort_fn(): done = True if rank == 0: print() print('Aborting...') # Save image snapshot. if (rank == 0) and (image_snapshot_ticks is not None) and (done or cur_tick % image_snapshot_ticks == 0): print('Saving images...') out = [] for image_app, image_mot, motion_app, motion, ref_tri, c, exp_vals in zip( images_app_val, images_mot_val, motion_1_out_vals, motion_2_out_vals, ref_tri_vals, grid_c_recon, cur_exp_out_vals): with torch.no_grad(): out.append( G_ema(image_app, image_mot, motion_app, motion, c=c, mesh= exp_vals[:, 2], triplane_recon=ref_tri, ws_avg=Rendering.ws_avg, motion_scale=motion_scale)) if 'image' in out[0]: images = torch.cat([o['image'].cpu() for o in out]).numpy() print(111111111111111111111111111111) print(images.shape) print(images.max()) print(images.min()) save_image_grid(images, os.path.join(run_dir, f'fakes{cur_nimg // 1000:06d}.png'), drange=[-1, 1], grid_size=grid_size) if 'image_depth' in out[0]: images_depth = -torch.cat([o['image_depth'].cpu() for o in out]).numpy() save_image_grid(images_depth, os.path.join(run_dir, f'fakes{cur_nimg // 1000:06d}_depth.png'), drange=[images_depth.min(), images_depth.max()], grid_size=grid_size) if 'image_sr' in out[0] and out[0]['image_sr'] is not None: images_sr = torch.cat([o['image_sr'].cpu() for o in out]).numpy() save_image_grid(images_sr, os.path.join(run_dir, f'fakes{cur_nimg // 1000:06d}_sr.png'), drange=[-1, 1], grid_size=grid_size) # Save network snapshot. snapshot_pkl = None snapshot_data = None if (network_snapshot_ticks is not None) and (done or cur_tick % network_snapshot_ticks == 0): snapshot_data = dict(training_set_kwargs=dict(training_set_kwargs)) for name, module in [('G', G), ('D', D), ('G_ema', G_ema), ('D_patch', D_patch), ('augment_pipe', augment_pipe)]: if module is not None: if num_gpus > 1: misc.check_ddp_consistency(module, ignore_regex=r'.*\.[^.]+_(avg|ema)') module = copy.deepcopy(module).eval().requires_grad_(False).cpu() snapshot_data[name] = module del module # conserve memory snapshot_pkl = os.path.join(run_dir, f'network-snapshot-{cur_nimg // 1000:06d}.pkl') if rank == 0: with open(snapshot_pkl, 'wb') as f: pickle.dump(snapshot_data, f) # Collect statistics. for phase in phases: value = [] if (phase.start_event is not None) and (phase.end_event is not None): phase.end_event.synchronize() value = phase.start_event.elapsed_time(phase.end_event) training_stats.report0('Timing/' + phase.name, value) stats_collector.update() stats_dict = stats_collector.as_dict() # Update logs. timestamp = time.time() if stats_jsonl is not None: fields = dict(stats_dict, timestamp=timestamp) stats_jsonl.write(json.dumps(fields) + '\n') stats_jsonl.flush() if stats_tfevents is not None: global_step = int(cur_nimg / 1e3) walltime = timestamp - start_time for name, value in stats_dict.items(): stats_tfevents.add_scalar(name, value.mean, global_step=global_step, walltime=walltime) for name, value in stats_metrics.items(): stats_tfevents.add_scalar(f'Metrics/{name}', value, global_step=global_step, walltime=walltime) stats_tfevents.flush() if progress_fn is not None: progress_fn(cur_nimg // 1000, total_kimg) # Update state. cur_tick += 1 tick_start_nimg = cur_nimg tick_start_time = time.time() maintenance_time = tick_start_time - tick_end_time if done: break # Done. if rank == 0: print() print('Exiting...') # ----------------------------------------------------------------------------