Spaces:
Build error
Build error
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| # Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. | |
| # | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| import imp | |
| import os | |
| import time | |
| import copy | |
| import json | |
| import pickle | |
| import psutil | |
| import PIL.Image | |
| import numpy as np | |
| import torch | |
| import dnnlib | |
| import tqdm | |
| import shutil | |
| import legacy | |
| 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 torch_utils.distributed_utils import gather_list_and_concat | |
| from metrics import metric_main | |
| from training.data_utils import save_image_grid, resize_image | |
| #---------------------------------------------------------------------------- | |
| def setup_snapshot_image_grid(training_set, random_seed=0): | |
| rnd = np.random.RandomState(random_seed) | |
| gw = np.clip(7680 // training_set.image_shape[2], 7, 32) | |
| gh = np.clip(4320 // training_set.image_shape[1], 4, 32) | |
| # No labels => show random subset of training samples. | |
| if not training_set.has_labels: | |
| all_indices = list(range(len(training_set))) | |
| rnd.shuffle(all_indices) | |
| grid_indices = [all_indices[i % len(all_indices)] for i in range(gw * gh)] | |
| else: | |
| # Group training samples by label. | |
| label_groups = dict() # label => [idx, ...] | |
| for idx in range(len(training_set)): | |
| label = tuple(training_set.get_details(idx).raw_label.flat[::-1]) | |
| if label not in label_groups: | |
| label_groups[label] = [] | |
| label_groups[label].append(idx) | |
| # Reorder. | |
| label_order = sorted(label_groups.keys()) | |
| for label in label_order: | |
| rnd.shuffle(label_groups[label]) | |
| # Organize into grid. | |
| grid_indices = [] | |
| for y in range(gh): | |
| label = label_order[y % len(label_order)] | |
| indices = label_groups[label] | |
| grid_indices += [indices[x % len(indices)] for x in range(gw)] | |
| label_groups[label] = [indices[(i + gw) % len(indices)] for i in range(len(indices))] | |
| # Load data. | |
| images, labels, _ = zip(*[training_set[i] for i in grid_indices]) | |
| return (gw, gh), np.stack(images), np.stack(labels) | |
| #---------------------------------------------------------------------------- | |
| 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. | |
| G_opt_kwargs = {}, # Options for generator optimizer. | |
| D_opt_kwargs = {}, # Options for discriminator optimizer. | |
| 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. | |
| world_size = 1, # Number of GPUs participating in the training. | |
| rank = 0, # Rank of the current process. | |
| gpu = 0, # Index of GPU used in training | |
| batch_gpu = 4, # Batch size for once GPU | |
| batch_size = 4, # Total batch size for one training iteration. Can be larger than batch_gpu * world_size. | |
| ema_kimg = 10, # Half-life of the exponential moving average (EMA) of generator weights. | |
| ema_rampup = None, # EMA ramp-up coefficient. | |
| G_reg_interval = 4, # 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. | |
| 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_start = 0, # Resume from steps | |
| cudnn_benchmark = True, # Enable torch.backends.cudnn.benchmark? | |
| allow_tf32 = False, # Enable torch.backends.cuda.matmul.allow_tf32 and torch.backends.cudnn.allow_tf32? | |
| 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. | |
| update_cam_prior_ticks = None, # (optional) Non-parameteric updating camera poses of the dataset | |
| generation_with_image = False, # (optional) For each random z, you also sample an image associated with it. | |
| **unused, | |
| ): | |
| # Initialize. | |
| start_time = time.time() | |
| device = torch.device('cuda', gpu) | |
| np.random.seed(random_seed * world_size + rank) | |
| torch.manual_seed(random_seed * world_size + rank) | |
| torch.backends.cudnn.benchmark = cudnn_benchmark # Improves training speed. | |
| torch.backends.cuda.matmul.allow_tf32 = allow_tf32 # Allow PyTorch to internally use tf32 for matmul | |
| torch.backends.cudnn.allow_tf32 = allow_tf32 # Allow PyTorch to internally use tf32 for convolutions | |
| conv2d_gradfix.enabled = True # Improves training speed. | |
| grid_sample_gradfix.enabled = True # Avoids errors with the augmentation pipe. | |
| img_dir = run_dir + '/images' | |
| os.makedirs(img_dir, exist_ok=True) | |
| assert batch_gpu <= (batch_size // world_size) | |
| # Load training set. | |
| if rank == 0: | |
| print('Loading training set...') | |
| if world_size == 1: | |
| data_loader_kwargs.update({'num_workers': 1, 'prefetch_factor': 1}) | |
| training_set = dnnlib.util.construct_class_by_name(**training_set_kwargs) # subclass of training.dataset.Dataset | |
| # Setup dataloader/sampler | |
| # if getattr(G.synthesis, 'sampler', None) is not None: | |
| # raise NotImplementedError('ERROR NEED TO TAKE A LOOK') | |
| # # training_set_sampler = G.synthesis.sampler[0]( | |
| # # dataset=training_set, rank=rank, num_replicas=world_size, | |
| # # seed=random_seed, device=device, **G.synthesis.sampler[1]) | |
| # else: | |
| training_set_sampler = misc.InfiniteSampler( | |
| dataset=training_set, rank=rank, num_replicas=world_size, seed=random_seed) | |
| training_set_iterator = iter(torch.utils.data.DataLoader( | |
| dataset=training_set, sampler=training_set_sampler, batch_size=batch_size//world_size, **data_loader_kwargs)) | |
| if generation_with_image: | |
| backup_data_iterator = iter(torch.utils.data.DataLoader( | |
| dataset=copy.deepcopy(training_set), sampler=training_set_sampler, batch_size=batch_size//world_size, **data_loader_kwargs)) | |
| if rank == 0: | |
| print() | |
| print('Num images: ', len(training_set)) | |
| print('Image shape:', training_set.image_shape) | |
| print('Label shape:', training_set.label_shape) | |
| print() | |
| # Construct networks. | |
| if rank == 0: | |
| print('Constructing networks...') | |
| common_kwargs = dict(c_dim=training_set.label_dim, img_resolution=training_set.resolution, img_channels=training_set.num_channels) | |
| if G_kwargs.get('img_channels', None) is not None: | |
| common_kwargs['img_channels'] = G_kwargs['img_channels'] | |
| del G_kwargs['img_channels'] | |
| G = dnnlib.util.construct_class_by_name(**G_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module | |
| if 'camera_kwargs' in D_kwargs: | |
| D_kwargs.camera_kwargs.w_dim, D_kwargs.camera_kwargs.num_ws = G.w_dim, G.num_ws | |
| D = dnnlib.util.construct_class_by_name(**D_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module | |
| G_ema = copy.deepcopy(G).eval() | |
| resize_real_img_early = D_kwargs.get('resize_real_early', False) | |
| disc_enable_ema = D_kwargs.get('enable_ema', False) | |
| if disc_enable_ema: | |
| D_ema = copy.deepcopy(D).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) | |
| modules = [('G', G), ('D', D), ('G_ema', G_ema)] | |
| if disc_enable_ema: | |
| modules += [('D_ema', D_ema)] | |
| for name, module in modules: | |
| misc.copy_params_and_buffers(resume_data[name], module, require_all=False) | |
| # Print network summary tables. | |
| if rank == 0: | |
| z = torch.empty([batch_gpu, G.z_dim], device=device) | |
| c = torch.empty([batch_gpu, G.c_dim], device=device) | |
| img = misc.print_module_summary(G, [z, c]) | |
| misc.print_module_summary(D, [img, c]) | |
| # 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 {world_size} GPUs...') | |
| ddp_modules = dict() | |
| module_list = [('G_mapping', G.mapping), ('G_synthesis', G.synthesis), ('D', D), (None, G_ema), ('augment_pipe', augment_pipe)] | |
| if G.encoder is not None: | |
| module_list += [('G_encoder', G.encoder)] | |
| if disc_enable_ema: | |
| module_list += [('D_ema', D_ema)] | |
| for name, module in module_list: | |
| if (world_size > 1) and (module is not None) and len(list(module.parameters())) != 0: | |
| module.requires_grad_(True) | |
| module = torch.nn.parallel.DistributedDataParallel( | |
| module, device_ids=[device], broadcast_buffers=False, find_unused_parameters=True) # allows progressive | |
| module.requires_grad_(False) | |
| if name is not None: | |
| ddp_modules[name] = module | |
| # Setup training phases. | |
| if rank == 0: | |
| print('Setting up training phases...') | |
| loss = dnnlib.util.construct_class_by_name(device=device, **ddp_modules, **loss_kwargs) # subclass of training.loss.Loss | |
| phases = [] | |
| for name, module, opt_kwargs, reg_interval in [('G', G, G_opt_kwargs, G_reg_interval), ('D', D, D_opt_kwargs, D_reg_interval)]: | |
| if reg_interval is None: | |
| opt = dnnlib.util.construct_class_by_name(params=module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer | |
| phases += [dnnlib.EasyDict(name=name+'both', module=module, opt=opt, interval=1, scaler=None)] | |
| 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(module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer | |
| phases += [dnnlib.EasyDict(name=name+'main', module=module, opt=opt, interval=1, scaler=None)] | |
| phases += [dnnlib.EasyDict(name=name+'reg', module=module, opt=opt, interval=reg_interval, scaler=None)] | |
| 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 | |
| grid_i = None | |
| if rank == 0: | |
| print(f'Exporting sample images... {batch_gpu}') | |
| grid_size, images, labels = setup_snapshot_image_grid(training_set=training_set) | |
| grid_z = torch.randn([labels.shape[0], G.z_dim], device=device).split(batch_gpu) | |
| grid_c = torch.from_numpy(labels).to(device).split(batch_gpu) | |
| grid_i = (torch.from_numpy(images).float() / 127.5 - 1).to(device).split(batch_gpu) | |
| if not os.path.exists(os.path.join(img_dir, 'reals.png')): | |
| save_image_grid(images, os.path.join(img_dir, 'reals.png'), drange=[0,255], grid_size=grid_size) | |
| if not os.path.exists( os.path.join(img_dir, 'fakes_init.png')): | |
| with torch.no_grad(): | |
| images = torch.cat([G_ema.get_final_output(z=z, c=c, noise_mode='const', img=img).cpu() for z, c, img in zip(grid_z, grid_c, grid_i)]).numpy() | |
| save_image_grid(images, os.path.join(img_dir, 'fakes_init.png'), drange=[-1,1], grid_size=grid_size) | |
| # 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'), 'at') | |
| 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_start | |
| cur_tick = cur_nimg // (1000 * kimg_per_tick) | |
| 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) | |
| while True: | |
| # set number of images | |
| loss.set_alpha(cur_nimg) | |
| curr_res = loss.resolution | |
| # Estimating Cameras for the training set (optional) | |
| if hasattr(training_set_sampler, 'update_dataset_cameras') and \ | |
| (cur_nimg == resume_start and resume_start > 0 and cur_tick > update_cam_prior_ticks): | |
| training_set_sampler.update_dataset_cameras(D.get_estimated_camera) | |
| # Fetch training data. | |
| with torch.autograd.profiler.record_function('data_fetch'): | |
| def load_data(iterator): | |
| img, c, _ = next(iterator) | |
| if resize_real_img_early: | |
| img = resize_image(img, curr_res) | |
| img = [{'img': img} for img in (img.to(device).to(torch.float32) / 127.5 - 1).split(batch_gpu)] | |
| c = c.to(device).split(batch_gpu) | |
| return img, c | |
| phase_real_img, phase_real_c = load_data(training_set_iterator) | |
| all_gen_z = torch.randn([len(phases) * batch_size, G.z_dim], device=device) | |
| all_gen_z = [phase_gen_z.split(batch_gpu) for phase_gen_z in all_gen_z.split(batch_size)] | |
| all_gen_c = [training_set.get_label(np.random.randint(len(training_set))) for _ in range(len(phases) * batch_size)] | |
| all_gen_c = torch.from_numpy(np.stack(all_gen_c)).pin_memory().to(device) | |
| all_gen_c = [phase_gen_c.split(batch_gpu) for phase_gen_c in all_gen_c.split(batch_size)] | |
| all_gen_img = [[None for _ in range(len(phase_real_img))] for _ in range(len(phases))] | |
| # Execute training phases. | |
| # with torch.autograd.profiler.profile(with_stack=True, profile_memory=True) as prof: | |
| for phase, phase_gen_z, phase_gen_c, phase_gen_img in zip(phases, all_gen_z, all_gen_c, all_gen_img): | |
| if batch_idx % phase.interval != 0: | |
| continue | |
| if generation_with_image: | |
| phase_gen_img, phase_gen_c = load_data(backup_data_iterator) | |
| # Initialize gradient accumulation. | |
| if phase.start_event is not None: | |
| phase.start_event.record(torch.cuda.current_stream(device)) | |
| phase.opt.zero_grad(set_to_none=True) | |
| phase.module.requires_grad_(True) | |
| # Accumulate gradients over multiple rounds. | |
| for round_idx, (real_img, real_c, gen_z, gen_c, fake_img) in enumerate(zip(phase_real_img, phase_real_c, phase_gen_z, phase_gen_c, phase_gen_img)): | |
| sync = (round_idx == batch_size // (batch_gpu * world_size) - 1) | |
| gain = phase.interval | |
| losses = loss.accumulate_gradients( | |
| phase=phase.name, | |
| real_img=real_img, | |
| real_c=real_c, gen_z=gen_z, | |
| gen_c=gen_c, fake_img=fake_img, | |
| sync=sync, gain=gain, scaler=phase.scaler) | |
| # Update weights. | |
| phase.module.requires_grad_(False) | |
| with torch.autograd.profiler.record_function(phase.name + '_opt'): | |
| if len(losses) > 0: | |
| if phase.scaler is not None: | |
| phase.scaler.unscale_(phase.opt) | |
| all_grads = [] | |
| for param in phase.module.parameters(): | |
| if param.grad is not None: | |
| misc.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) | |
| all_grads += [torch.norm(param.grad.detach(), p=2)] | |
| grad_norm = torch.stack(all_grads).norm(p=2) | |
| if phase.scaler is not None: | |
| phase.scaler.step(phase.opt) | |
| phase.scaler.update() | |
| training_stats.report(f'Scaler/{phase.name}', phase.scaler.get_scale()) | |
| else: | |
| phase.opt.step() | |
| training_stats.report(f'Gradient/{phase.name}', grad_norm) | |
| 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) | |
| if disc_enable_ema: # update EMA for discriminator | |
| for p_ema, p in zip(D_ema.parameters(), D.parameters()): | |
| p_ema.copy_(p.lerp(p_ema, ema_beta)) | |
| for b_ema, b in zip(D_ema.buffers(), D.buffers()): | |
| b_ema.copy_(b) | |
| # 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 stats_collector. | |
| tick_end_time = time.time() | |
| fields = [f"[{run_dir}]:"] | |
| fields += [f"tick {training_stats.report0('Progress/tick', cur_tick):<5d}"] | |
| fields += [f"kimg {training_stats.report0('Progress/kimg', cur_nimg / 1e3):<8.1f}"] | |
| if loss.alpha is not None: | |
| fields += [f"alpha {training_stats.report0('Progress/alpha', loss.alpha):<8.5f}"] | |
| fields += [f"res {training_stats.report0('Progress/res', loss.resolution):<5d}"] | |
| 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/maintain_sec', maintenance_time):<5.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}"] | |
| 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...') | |
| # Update the dataset cameras? | |
| if hasattr(training_set_sampler, 'update_dataset_cameras') and \ | |
| (cur_tick % update_cam_prior_ticks == 0 and cur_tick > 0): | |
| training_set_sampler.update_dataset_cameras(D.get_estimated_camera) | |
| # Save image snapshot. | |
| if (rank == 0) and (image_snapshot_ticks is not None) and (done or cur_tick % image_snapshot_ticks == 0): | |
| with torch.no_grad(): | |
| images = torch.cat([G_ema.get_final_output(z=z, c=c, noise_mode='const', img=None).cpu() for z, c, img in zip(grid_z, grid_c, grid_i)]).numpy() | |
| save_image_grid(images, os.path.join(img_dir, f'fakes{cur_nimg//1000:06d}.png'), drange=[-1,1], grid_size=grid_size) | |
| images = torch.cat([G_ema.get_final_output(z=z, c=c, noise_mode='const', img=img, camera_mode=[0.5,0.5,0.5]).cpu() for z, c, img in zip(grid_z, grid_c, grid_i)]).numpy() | |
| save_image_grid(images, os.path.join(img_dir, f'fakes{cur_nimg//1000:06d}_000.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)) | |
| modules = [('G', G), ('D', D), ('G_ema', G_ema), ('augment_pipe', augment_pipe)] | |
| if disc_enable_ema: | |
| modules += [('D_ema', D_ema)] | |
| for name, module in modules: | |
| if module is not None: | |
| 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) | |
| # save the latest checkpoint | |
| shutil.copy(snapshot_pkl, os.path.join(run_dir, 'latest-network-snapshot.pkl')) | |
| # Evaluate metrics. | |
| if (snapshot_data is not None) and (len(metrics) > 0) and (cur_tick > 1): | |
| if rank == 0: | |
| print('Evaluating metrics...') | |
| for metric in metrics: | |
| result_dict = metric_main.calc_metric(metric=metric, G=snapshot_data['G_ema'], | |
| dataset_kwargs=training_set_kwargs, num_gpus=world_size, rank=rank, device=device) | |
| if rank == 0: | |
| metric_main.report_metric(result_dict, run_dir=run_dir, snapshot_pkl=snapshot_pkl) | |
| stats_metrics.update(result_dict.results) | |
| del snapshot_data # conserve memory | |
| # 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 rank == 0: | |
| losses = [(key, fields[key]) for key in fields if 'Loss/' in key] | |
| losses = ["{}: {:.4f}".format(key[5:], loss['mean']) for key, loss in losses] | |
| print('\t'.join(losses)) | |
| 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...') | |
| #---------------------------------------------------------------------------- | |