File size: 15,506 Bytes
3494c6b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 |
"""
Utilities for instrumenting a torch model.
Trace will hook one layer at a time.
TraceDict will hook multiple layers at once.
subsequence slices intervals from Sequential modules.
get_module, replace_module, get_parameter resolve dotted names.
set_requires_grad recursively sets requires_grad in module parameters.
"""
import contextlib
import copy
import inspect
from collections import OrderedDict
import torch
class Trace(contextlib.AbstractContextManager):
"""
To retain the output of the named layer during the computation of
the given network:
with Trace(net, 'layer.name') as ret:
_ = net(inp)
representation = ret.output
A layer module can be passed directly without a layer name, and
its output will be retained. By default, a direct reference to
the output object is returned, but options can control this:
clone=True - retains a copy of the output, which can be
useful if you want to see the output before it might
be modified by the network in-place later.
detach=True - retains a detached reference or copy. (By
default the value would be left attached to the graph.)
retain_grad=True - request gradient to be retained on the
output. After backward(), ret.output.grad is populated.
retain_input=True - also retains the input.
retain_output=False - can disable retaining the output.
edit_output=fn - calls the function to modify the output
of the layer before passing it the rest of the model.
fn can optionally accept (output, layer) arguments
for the original output and the layer name.
stop=True - throws a StopForward exception after the layer
is run, which allows running just a portion of a model.
"""
def __init__(
self,
module,
layer=None,
retain_output=True,
retain_input=False,
clone=False,
detach=False,
retain_grad=False,
edit_output=None,
stop=False,
):
"""
Method to replace a forward method with a closure that
intercepts the call, and tracks the hook so that it can be reverted.
"""
retainer = self
self.layer = layer
if layer is not None:
module = get_module(module, layer)
def retain_hook(m, inputs, output):
if retain_input:
retainer.input = recursive_copy(
inputs[0] if len(inputs) == 1 else inputs,
clone=clone,
detach=detach,
retain_grad=False,
) # retain_grad applies to output only.
if edit_output:
output = invoke_with_optional_args(
edit_output, output=output, layer=self.layer
)
if retain_output:
retainer.output = recursive_copy(
output, clone=clone, detach=detach, retain_grad=retain_grad
)
# When retain_grad is set, also insert a trivial
# copy operation. That allows in-place operations
# to follow without error.
if retain_grad:
output = recursive_copy(retainer.output, clone=True, detach=False)
if stop:
raise StopForward()
return output
self.registered_hook = module.register_forward_hook(retain_hook)
self.stop = stop
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
if self.stop and issubclass(type, StopForward):
return True
def close(self):
self.registered_hook.remove()
class TraceDict(OrderedDict, contextlib.AbstractContextManager):
"""
To retain the output of multiple named layers during the computation
of the given network:
with TraceDict(net, ['layer1.name1', 'layer2.name2']) as ret:
_ = net(inp)
representation = ret['layer1.name1'].output
If edit_output is provided, it should be a function that takes
two arguments: output, and the layer name; and then it returns the
modified output.
Other arguments are the same as Trace. If stop is True, then the
execution of the network will be stopped after the last layer
listed (even if it would not have been the last to be executed).
"""
def __init__(
self,
module,
layers=None,
retain_output=True,
retain_input=False,
clone=False,
detach=False,
retain_grad=False,
edit_output=None,
stop=False,
):
self.stop = stop
def flag_last_unseen(it):
try:
it = iter(it)
prev = next(it)
seen = set([prev])
except StopIteration:
return
for item in it:
if item not in seen:
yield False, prev
seen.add(item)
prev = item
yield True, prev
for is_last, layer in flag_last_unseen(layers):
self[layer] = Trace(
module=module,
layer=layer,
retain_output=retain_output,
retain_input=retain_input,
clone=clone,
detach=detach,
retain_grad=retain_grad,
edit_output=edit_output,
stop=stop and is_last,
)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
if self.stop and issubclass(type, StopForward):
return True
def close(self):
for layer, trace in reversed(self.items()):
trace.close()
class StopForward(Exception):
"""
If the only output needed from running a network is the retained
submodule then Trace(submodule, stop=True) will stop execution
immediately after the retained submodule by raising the StopForward()
exception. When Trace is used as context manager, it catches that
exception and can be used as follows:
with Trace(net, layername, stop=True) as tr:
net(inp) # Only runs the network up to layername
print(tr.output)
"""
pass
def recursive_copy(x, clone=None, detach=None, retain_grad=None):
"""
Copies a reference to a tensor, or an object that contains tensors,
optionally detaching and cloning the tensor(s). If retain_grad is
true, the original tensors are marked to have grads retained.
"""
if not clone and not detach and not retain_grad:
return x
if isinstance(x, torch.Tensor):
if retain_grad:
if not x.requires_grad:
x.requires_grad = True
x.retain_grad()
elif detach:
x = x.detach()
if clone:
x = x.clone()
return x
# Only dicts, lists, and tuples (and subclasses) can be copied.
if isinstance(x, dict):
return type(x)({k: recursive_copy(v) for k, v in x.items()})
elif isinstance(x, (list, tuple)):
return type(x)([recursive_copy(v) for v in x])
else:
assert False, f"Unknown type {type(x)} cannot be broken into tensors."
def subsequence(
sequential,
first_layer=None,
last_layer=None,
after_layer=None,
upto_layer=None,
single_layer=None,
share_weights=False,
):
"""
Creates a subsequence of a pytorch Sequential model, copying over
modules together with parameters for the subsequence. Only
modules from first_layer to last_layer (inclusive) are included,
or modules between after_layer and upto_layer (exclusive).
Handles descent into dotted layer names as long as all references
are within nested Sequential models.
If share_weights is True, then references the original modules
and their parameters without copying them. Otherwise, by default,
makes a separate brand-new copy.
"""
assert (single_layer is None) or (
first_layer is last_layer is after_layer is upto_layer is None
)
if single_layer is not None:
first_layer = single_layer
last_layer = single_layer
first, last, after, upto = [
None if d is None else d.split(".")
for d in [first_layer, last_layer, after_layer, upto_layer]
]
return hierarchical_subsequence(
sequential,
first=first,
last=last,
after=after,
upto=upto,
share_weights=share_weights,
)
def hierarchical_subsequence(
sequential, first, last, after, upto, share_weights=False, depth=0
):
"""
Recursive helper for subsequence() to support descent into dotted
layer names. In this helper, first, last, after, and upto are
arrays of names resulting from splitting on dots. Can only
descend into nested Sequentials.
"""
assert (last is None) or (upto is None)
assert (first is None) or (after is None)
if first is last is after is upto is None:
return sequential if share_weights else copy.deepcopy(sequential)
assert isinstance(sequential, torch.nn.Sequential), (
".".join((first or last or after or upto)[:depth] or "arg") + " not Sequential"
)
including_children = (first is None) and (after is None)
included_children = OrderedDict()
# A = current level short name of A.
# AN = full name for recursive descent if not innermost.
(F, FN), (L, LN), (A, AN), (U, UN) = [
(d[depth], (None if len(d) == depth + 1 else d))
if d is not None
else (None, None)
for d in [first, last, after, upto]
]
for name, layer in sequential._modules.items():
if name == F:
first = None
including_children = True
if name == A and AN is not None: # just like F if not a leaf.
after = None
including_children = True
if name == U and UN is None:
upto = None
including_children = False
if including_children:
# AR = full name for recursive descent if name matches.
FR, LR, AR, UR = [
n if n is None or n[depth] == name else None for n in [FN, LN, AN, UN]
]
chosen = hierarchical_subsequence(
layer,
first=FR,
last=LR,
after=AR,
upto=UR,
share_weights=share_weights,
depth=depth + 1,
)
if chosen is not None:
included_children[name] = chosen
if name == L:
last = None
including_children = False
if name == U and UN is not None: # just like L if not a leaf.
upto = None
including_children = False
if name == A and AN is None:
after = None
including_children = True
for name in [first, last, after, upto]:
if name is not None:
raise ValueError("Layer %s not found" % ".".join(name))
# Omit empty subsequences except at the outermost level,
# where we should not return None.
if not len(included_children) and depth > 0:
return None
result = torch.nn.Sequential(included_children)
result.training = sequential.training
return result
def set_requires_grad(requires_grad, *models):
"""
Sets requires_grad true or false for all parameters within the
models passed.
"""
for model in models:
if isinstance(model, torch.nn.Module):
for param in model.parameters():
param.requires_grad = requires_grad
elif isinstance(model, (torch.nn.Parameter, torch.Tensor)):
model.requires_grad = requires_grad
else:
assert False, "unknown type %r" % type(model)
def get_module(model, name):
"""
Finds the named module within the given model.
"""
for n, m in model.named_modules():
if n == name:
return m
raise LookupError(name)
def get_parameter(model, name):
"""
Finds the named parameter within the given model.
"""
for n, p in model.named_parameters():
if n == name:
return p
raise LookupError(name)
def replace_module(model, name, new_module):
"""
Replaces the named module within the given model.
"""
if "." in name:
parent_name, attr_name = name.rsplit(".", 1)
model = get_module(model, parent_name)
# original_module = getattr(model, attr_name)
setattr(model, attr_name, new_module)
def invoke_with_optional_args(fn, *args, **kwargs):
"""
Invokes a function with only the arguments that it
is written to accept, giving priority to arguments
that match by-name, using the following rules.
(1) arguments with matching names are passed by name.
(2) remaining non-name-matched args are passed by order.
(3) extra caller arguments that the function cannot
accept are not passed.
(4) extra required function arguments that the caller
cannot provide cause a TypeError to be raised.
Ordinary python calling conventions are helpful for
supporting a function that might be revised to accept
extra arguments in a newer version, without requiring the
caller to pass those new arguments. This function helps
support function callers that might be revised to supply
extra arguments, without requiring the callee to accept
those new arguments.
"""
argspec = inspect.getfullargspec(fn)
pass_args = []
used_kw = set()
unmatched_pos = []
used_pos = 0
defaulted_pos = len(argspec.args) - (
0 if not argspec.defaults else len(argspec.defaults)
)
# Pass positional args that match name first, then by position.
for i, n in enumerate(argspec.args):
if n in kwargs:
pass_args.append(kwargs[n])
used_kw.add(n)
elif used_pos < len(args):
pass_args.append(args[used_pos])
used_pos += 1
else:
unmatched_pos.append(len(pass_args))
pass_args.append(
None if i < defaulted_pos else argspec.defaults[i - defaulted_pos]
)
# Fill unmatched positional args with unmatched keyword args in order.
if len(unmatched_pos):
for k, v in kwargs.items():
if k in used_kw or k in argspec.kwonlyargs:
continue
pass_args[unmatched_pos[0]] = v
used_kw.add(k)
unmatched_pos = unmatched_pos[1:]
if len(unmatched_pos) == 0:
break
else:
if unmatched_pos[0] < defaulted_pos:
unpassed = ", ".join(
argspec.args[u] for u in unmatched_pos if u < defaulted_pos
)
raise TypeError(f"{fn.__name__}() cannot be passed {unpassed}.")
# Pass remaining kw args if they can be accepted.
pass_kw = {
k: v
for k, v in kwargs.items()
if k not in used_kw and (k in argspec.kwonlyargs or argspec.varargs is not None)
}
# Pass remaining positional args if they can be accepted.
if argspec.varargs is not None:
pass_args += list(args[used_pos:])
return fn(*pass_args, **pass_kw)
|