riciii7/FastAPI-Batik-GAN
0
1# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto. Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9import re10import contextlib11import numpy as np12import torch13import warnings14import dnnlib15 16#----------------------------------------------------------------------------17# Cached construction of constant tensors. Avoids CPU=>GPU copy when the18# same constant is used multiple times.19 20_constant_cache = dict()21 22def constant(value, shape=None, dtype=None, device=None, memory_format=None):23 value = np.asarray(value)24 if shape is not None:25 shape = tuple(shape)26 if dtype is None:27 dtype = torch.get_default_dtype()28 if device is None:29 device = torch.device('cpu')30 if memory_format is None:31 memory_format = torch.contiguous_format32 33 key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)34 tensor = _constant_cache.get(key, None)35 if tensor is None:36 tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)37 if shape is not None:38 tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))39 tensor = tensor.contiguous(memory_format=memory_format)40 _constant_cache[key] = tensor41 return tensor42 43#----------------------------------------------------------------------------44# Replace NaN/Inf with specified numerical values.45 46try:47 nan_to_num = torch.nan_to_num # 1.8.0a048except AttributeError:49 def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin50 assert isinstance(input, torch.Tensor)51 if posinf is None:52 posinf = torch.finfo(input.dtype).max53 if neginf is None:54 neginf = torch.finfo(input.dtype).min55 assert nan == 056 return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)57 58#----------------------------------------------------------------------------59# Symbolic assert.60 61try:62 symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access63except AttributeError:64 symbolic_assert = torch.Assert # 1.7.065 66#----------------------------------------------------------------------------67# Context manager to suppress known warnings in torch.jit.trace().68 69class suppress_tracer_warnings(warnings.catch_warnings):70 def __enter__(self):71 super().__enter__()72 warnings.simplefilter('ignore', category=torch.jit.TracerWarning)73 return self74 75#----------------------------------------------------------------------------76# Assert that the shape of a tensor matches the given list of integers.77# None indicates that the size of a dimension is allowed to vary.78# Performs symbolic assertion when used in torch.jit.trace().79 80def assert_shape(tensor, ref_shape):81 if tensor.ndim != len(ref_shape):82 raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')83 for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):84 if ref_size is None:85 pass86 elif isinstance(ref_size, torch.Tensor):87 with suppress_tracer_warnings(): # as_tensor results are registered as constants88 symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')89 elif isinstance(size, torch.Tensor):90 with suppress_tracer_warnings(): # as_tensor results are registered as constants91 symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')92 elif size != ref_size:93 raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')94 95#----------------------------------------------------------------------------96# Function decorator that calls torch.autograd.profiler.record_function().97 98def profiled_function(fn):99 def decorator(*args, **kwargs):100 with torch.autograd.profiler.record_function(fn.__name__):101 return fn(*args, **kwargs)102 decorator.__name__ = fn.__name__103 return decorator104 105#----------------------------------------------------------------------------106# Sampler for torch.utils.data.DataLoader that loops over the dataset107# indefinitely, shuffling items as it goes.108 109class InfiniteSampler(torch.utils.data.Sampler):110 def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5):111 assert len(dataset) > 0112 assert num_replicas > 0113 assert 0 <= rank < num_replicas114 assert 0 <= window_size <= 1115 super().__init__(dataset)116 self.dataset = dataset117 self.rank = rank118 self.num_replicas = num_replicas119 self.shuffle = shuffle120 self.seed = seed121 self.window_size = window_size122 123 def __iter__(self):124 order = np.arange(len(self.dataset))125 rnd = None126 window = 0127 if self.shuffle:128 rnd = np.random.RandomState(self.seed)129 rnd.shuffle(order)130 window = int(np.rint(order.size * self.window_size))131 132 idx = 0133 while True:134 i = idx % order.size135 if idx % self.num_replicas == self.rank:136 yield order[i]137 if window >= 2:138 j = (i - rnd.randint(window)) % order.size139 order[i], order[j] = order[j], order[i]140 idx += 1141 142#----------------------------------------------------------------------------143# Utilities for operating with torch.nn.Module parameters and buffers.144 145def params_and_buffers(module):146 assert isinstance(module, torch.nn.Module)147 return list(module.parameters()) + list(module.buffers())148 149def named_params_and_buffers(module):150 assert isinstance(module, torch.nn.Module)151 return list(module.named_parameters()) + list(module.named_buffers())152 153def copy_params_and_buffers(src_module, dst_module, require_all=False):154 assert isinstance(src_module, torch.nn.Module)155 assert isinstance(dst_module, torch.nn.Module)156 src_tensors = {name: tensor for name, tensor in named_params_and_buffers(src_module)}157 for name, tensor in named_params_and_buffers(dst_module):158 assert (name in src_tensors) or (not require_all)159 if name in src_tensors:160 tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)161 162#----------------------------------------------------------------------------163# Context manager for easily enabling/disabling DistributedDataParallel164# synchronization.165 166@contextlib.contextmanager167def ddp_sync(module, sync):168 assert isinstance(module, torch.nn.Module)169 if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):170 yield171 else:172 with module.no_sync():173 yield174 175#----------------------------------------------------------------------------176# Check DistributedDataParallel consistency across processes.177 178def check_ddp_consistency(module, ignore_regex=None):179 assert isinstance(module, torch.nn.Module)180 for name, tensor in named_params_and_buffers(module):181 fullname = type(module).__name__ + '.' + name182 if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):183 continue184 tensor = tensor.detach()185 other = tensor.clone()186 torch.distributed.broadcast(tensor=other, src=0)187 assert (nan_to_num(tensor) == nan_to_num(other)).all(), fullname188 189#----------------------------------------------------------------------------190# Print summary table of module hierarchy.191 192def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):193 assert isinstance(module, torch.nn.Module)194 assert not isinstance(module, torch.jit.ScriptModule)195 assert isinstance(inputs, (tuple, list))196 197 # Register hooks.198 entries = []199 nesting = [0]200 def pre_hook(_mod, _inputs):201 nesting[0] += 1202 def post_hook(mod, _inputs, outputs):203 nesting[0] -= 1204 if nesting[0] <= max_nesting:205 outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]206 outputs = [t for t in outputs if isinstance(t, torch.Tensor)]207 entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs))208 hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]209 hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]210 211 # Run module.212 outputs = module(*inputs)213 for hook in hooks:214 hook.remove()215 216 # Identify unique outputs, parameters, and buffers.217 tensors_seen = set()218 for e in entries:219 e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]220 e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]221 e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]222 tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}223 224 # Filter out redundant entries.225 if skip_redundant:226 entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]227 228 # Construct table.229 rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]230 rows += [['---'] * len(rows[0])]231 param_total = 0232 buffer_total = 0233 submodule_names = {mod: name for name, mod in module.named_modules()}234 for e in entries:235 name = '<top-level>' if e.mod is module else submodule_names[e.mod]236 param_size = sum(t.numel() for t in e.unique_params)237 buffer_size = sum(t.numel() for t in e.unique_buffers)238 output_shapes = [str(list(e.outputs[0].shape)) for t in e.outputs]239 output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]240 rows += [[241 name + (':0' if len(e.outputs) >= 2 else ''),242 str(param_size) if param_size else '-',243 str(buffer_size) if buffer_size else '-',244 (output_shapes + ['-'])[0],245 (output_dtypes + ['-'])[0],246 ]]247 for idx in range(1, len(e.outputs)):248 rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]249 param_total += param_size250 buffer_total += buffer_size251 rows += [['---'] * len(rows[0])]252 rows += [['Total', str(param_total), str(buffer_total), '-', '-']]253 254 # Print table.255 widths = [max(len(cell) for cell in column) for column in zip(*rows)]256 print()257 for row in rows:258 print(' '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))259 print()260 return outputs261 262#----------------------------------------------------------------------------263