lambda/generative-music-visualizer

1# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  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 temporarily suppress known warnings in torch.jit.trace().68# Note: Cannot use catch_warnings because of https://bugs.python.org/issue2967269 70@contextlib.contextmanager71def suppress_tracer_warnings():72    flt = ('ignore', None, torch.jit.TracerWarning, None, 0)73    warnings.filters.insert(0, flt)74    yield75    warnings.filters.remove(flt)76 77#----------------------------------------------------------------------------78# Assert that the shape of a tensor matches the given list of integers.79# None indicates that the size of a dimension is allowed to vary.80# Performs symbolic assertion when used in torch.jit.trace().81 82def assert_shape(tensor, ref_shape):83    if tensor.ndim != len(ref_shape):84        raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')85    for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):86        if ref_size is None:87            pass88        elif isinstance(ref_size, torch.Tensor):89            with suppress_tracer_warnings(): # as_tensor results are registered as constants90                symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')91        elif isinstance(size, torch.Tensor):92            with suppress_tracer_warnings(): # as_tensor results are registered as constants93                symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')94        elif size != ref_size:95            raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')96 97#----------------------------------------------------------------------------98# Function decorator that calls torch.autograd.profiler.record_function().99 100def profiled_function(fn):101    def decorator(*args, **kwargs):102        with torch.autograd.profiler.record_function(fn.__name__):103            return fn(*args, **kwargs)104    decorator.__name__ = fn.__name__105    return decorator106 107#----------------------------------------------------------------------------108# Sampler for torch.utils.data.DataLoader that loops over the dataset109# indefinitely, shuffling items as it goes.110 111class InfiniteSampler(torch.utils.data.Sampler):112    def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5):113        assert len(dataset) > 0114        assert num_replicas > 0115        assert 0 <= rank < num_replicas116        assert 0 <= window_size <= 1117        super().__init__(dataset)118        self.dataset = dataset119        self.rank = rank120        self.num_replicas = num_replicas121        self.shuffle = shuffle122        self.seed = seed123        self.window_size = window_size124 125    def __iter__(self):126        order = np.arange(len(self.dataset))127        rnd = None128        window = 0129        if self.shuffle:130            rnd = np.random.RandomState(self.seed)131            rnd.shuffle(order)132            window = int(np.rint(order.size * self.window_size))133 134        idx = 0135        while True:136            i = idx % order.size137            if idx % self.num_replicas == self.rank:138                yield order[i]139            if window >= 2:140                j = (i - rnd.randint(window)) % order.size141                order[i], order[j] = order[j], order[i]142            idx += 1143 144#----------------------------------------------------------------------------145# Utilities for operating with torch.nn.Module parameters and buffers.146 147def params_and_buffers(module):148    assert isinstance(module, torch.nn.Module)149    return list(module.parameters()) + list(module.buffers())150 151def named_params_and_buffers(module):152    assert isinstance(module, torch.nn.Module)153    return list(module.named_parameters()) + list(module.named_buffers())154 155def copy_params_and_buffers(src_module, dst_module, require_all=False):156    assert isinstance(src_module, torch.nn.Module)157    assert isinstance(dst_module, torch.nn.Module)158    src_tensors = dict(named_params_and_buffers(src_module))159    for name, tensor in named_params_and_buffers(dst_module):160        assert (name in src_tensors) or (not require_all)161        if name in src_tensors:162            tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)163 164#----------------------------------------------------------------------------165# Context manager for easily enabling/disabling DistributedDataParallel166# synchronization.167 168@contextlib.contextmanager169def ddp_sync(module, sync):170    assert isinstance(module, torch.nn.Module)171    if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):172        yield173    else:174        with module.no_sync():175            yield176 177#----------------------------------------------------------------------------178# Check DistributedDataParallel consistency across processes.179 180def check_ddp_consistency(module, ignore_regex=None):181    assert isinstance(module, torch.nn.Module)182    for name, tensor in named_params_and_buffers(module):183        fullname = type(module).__name__ + '.' + name184        if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):185            continue186        tensor = tensor.detach()187        if tensor.is_floating_point():188            tensor = nan_to_num(tensor)189        other = tensor.clone()190        torch.distributed.broadcast(tensor=other, src=0)191        assert (tensor == other).all(), fullname192 193#----------------------------------------------------------------------------194# Print summary table of module hierarchy.195 196def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):197    assert isinstance(module, torch.nn.Module)198    assert not isinstance(module, torch.jit.ScriptModule)199    assert isinstance(inputs, (tuple, list))200 201    # Register hooks.202    entries = []203    nesting = [0]204    def pre_hook(_mod, _inputs):205        nesting[0] += 1206    def post_hook(mod, _inputs, outputs):207        nesting[0] -= 1208        if nesting[0] <= max_nesting:209            outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]210            outputs = [t for t in outputs if isinstance(t, torch.Tensor)]211            entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs))212    hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]213    hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]214 215    # Run module.216    outputs = module(*inputs)217    for hook in hooks:218        hook.remove()219 220    # Identify unique outputs, parameters, and buffers.221    tensors_seen = set()222    for e in entries:223        e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]224        e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]225        e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]226        tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}227 228    # Filter out redundant entries.229    if skip_redundant:230        entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]231 232    # Construct table.233    rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]234    rows += [['---'] * len(rows[0])]235    param_total = 0236    buffer_total = 0237    submodule_names = {mod: name for name, mod in module.named_modules()}238    for e in entries:239        name = '<top-level>' if e.mod is module else submodule_names[e.mod]240        param_size = sum(t.numel() for t in e.unique_params)241        buffer_size = sum(t.numel() for t in e.unique_buffers)242        output_shapes = [str(list(t.shape)) for t in e.outputs]243        output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]244        rows += [[245            name + (':0' if len(e.outputs) >= 2 else ''),246            str(param_size) if param_size else '-',247            str(buffer_size) if buffer_size else '-',248            (output_shapes + ['-'])[0],249            (output_dtypes + ['-'])[0],250        ]]251        for idx in range(1, len(e.outputs)):252            rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]253        param_total += param_size254        buffer_total += buffer_size255    rows += [['---'] * len(rows[0])]256    rows += [['Total', str(param_total), str(buffer_total), '-', '-']]257 258    # Print table.259    widths = [max(len(cell) for cell in column) for column in zip(*rows)]260    print()261    for row in rows:262        print('  '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))263    print()264    return outputs265 266#----------------------------------------------------------------------------267