CorvaeOboro commited on
Commit
e91c77c
β€’
1 Parent(s): debdde2

Upload misc.py

Browse files
Files changed (1) hide show
  1. torch_utils/misc.py +262 -0
torch_utils/misc.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ο»Ώ# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ import re
10
+ import contextlib
11
+ import numpy as np
12
+ import torch
13
+ import warnings
14
+ import dnnlib
15
+
16
+ #----------------------------------------------------------------------------
17
+ # Cached construction of constant tensors. Avoids CPU=>GPU copy when the
18
+ # same constant is used multiple times.
19
+
20
+ _constant_cache = dict()
21
+
22
+ def 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_format
32
+
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] = tensor
41
+ return tensor
42
+
43
+ #----------------------------------------------------------------------------
44
+ # Replace NaN/Inf with specified numerical values.
45
+
46
+ try:
47
+ nan_to_num = torch.nan_to_num # 1.8.0a0
48
+ except AttributeError:
49
+ def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin
50
+ assert isinstance(input, torch.Tensor)
51
+ if posinf is None:
52
+ posinf = torch.finfo(input.dtype).max
53
+ if neginf is None:
54
+ neginf = torch.finfo(input.dtype).min
55
+ assert nan == 0
56
+ return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)
57
+
58
+ #----------------------------------------------------------------------------
59
+ # Symbolic assert.
60
+
61
+ try:
62
+ symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access
63
+ except AttributeError:
64
+ symbolic_assert = torch.Assert # 1.7.0
65
+
66
+ #----------------------------------------------------------------------------
67
+ # Context manager to suppress known warnings in torch.jit.trace().
68
+
69
+ class suppress_tracer_warnings(warnings.catch_warnings):
70
+ def __enter__(self):
71
+ super().__enter__()
72
+ warnings.simplefilter('ignore', category=torch.jit.TracerWarning)
73
+ return self
74
+
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
+
80
+ def 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
+ pass
86
+ elif isinstance(ref_size, torch.Tensor):
87
+ with suppress_tracer_warnings(): # as_tensor results are registered as constants
88
+ 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 constants
91
+ 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
+
98
+ def 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 decorator
104
+
105
+ #----------------------------------------------------------------------------
106
+ # Sampler for torch.utils.data.DataLoader that loops over the dataset
107
+ # indefinitely, shuffling items as it goes.
108
+
109
+ class 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) > 0
112
+ assert num_replicas > 0
113
+ assert 0 <= rank < num_replicas
114
+ assert 0 <= window_size <= 1
115
+ super().__init__(dataset)
116
+ self.dataset = dataset
117
+ self.rank = rank
118
+ self.num_replicas = num_replicas
119
+ self.shuffle = shuffle
120
+ self.seed = seed
121
+ self.window_size = window_size
122
+
123
+ def __iter__(self):
124
+ order = np.arange(len(self.dataset))
125
+ rnd = None
126
+ window = 0
127
+ 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 = 0
133
+ while True:
134
+ i = idx % order.size
135
+ if idx % self.num_replicas == self.rank:
136
+ yield order[i]
137
+ if window >= 2:
138
+ j = (i - rnd.randint(window)) % order.size
139
+ order[i], order[j] = order[j], order[i]
140
+ idx += 1
141
+
142
+ #----------------------------------------------------------------------------
143
+ # Utilities for operating with torch.nn.Module parameters and buffers.
144
+
145
+ def params_and_buffers(module):
146
+ assert isinstance(module, torch.nn.Module)
147
+ return list(module.parameters()) + list(module.buffers())
148
+
149
+ def named_params_and_buffers(module):
150
+ assert isinstance(module, torch.nn.Module)
151
+ return list(module.named_parameters()) + list(module.named_buffers())
152
+
153
+ def 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 DistributedDataParallel
164
+ # synchronization.
165
+
166
+ @contextlib.contextmanager
167
+ def ddp_sync(module, sync):
168
+ assert isinstance(module, torch.nn.Module)
169
+ if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):
170
+ yield
171
+ else:
172
+ with module.no_sync():
173
+ yield
174
+
175
+ #----------------------------------------------------------------------------
176
+ # Check DistributedDataParallel consistency across processes.
177
+
178
+ def 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__ + '.' + name
182
+ if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):
183
+ continue
184
+ 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(), fullname
188
+
189
+ #----------------------------------------------------------------------------
190
+ # Print summary table of module hierarchy.
191
+
192
+ def 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] += 1
202
+ def post_hook(mod, _inputs, outputs):
203
+ nesting[0] -= 1
204
+ 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 = 0
232
+ buffer_total = 0
233
+ 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_size
250
+ buffer_total += buffer_size
251
+ 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 outputs
261
+
262
+ #----------------------------------------------------------------------------