| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354 |
- # mypy: allow-untyped-defs
- import inspect
- import torch
- def skip_init(module_cls, *args, **kwargs):
- r"""
- Given a module class object and args / kwargs, instantiate the module without initializing parameters / buffers.
- This can be useful if initialization is slow or if custom initialization will
- be performed, making the default initialization unnecessary. There are some caveats to this, due to
- the way this function is implemented:
- 1. The module must accept a `device` arg in its constructor that is passed to any parameters
- or buffers created during construction.
- 2. The module must not perform any computation on parameters in its constructor except
- initialization (i.e. functions from :mod:`torch.nn.init`).
- If these conditions are satisfied, the module can be instantiated with parameter / buffer values
- uninitialized, as if having been created using :func:`torch.empty`.
- Args:
- module_cls: Class object; should be a subclass of :class:`torch.nn.Module`
- args: args to pass to the module's constructor
- kwargs: kwargs to pass to the module's constructor
- Returns:
- Instantiated module with uninitialized parameters / buffers
- Example::
- >>> # xdoctest: +IGNORE_WANT("non-deterministic")
- >>> import torch
- >>> m = torch.nn.utils.skip_init(torch.nn.Linear, 5, 1)
- >>> m.weight
- Parameter containing:
- tensor([[0.0000e+00, 1.5846e+29, 7.8307e+00, 2.5250e-29, 1.1210e-44]],
- requires_grad=True)
- >>> m2 = torch.nn.utils.skip_init(torch.nn.Linear, in_features=6, out_features=1)
- >>> m2.weight
- Parameter containing:
- tensor([[-1.4677e+24, 4.5915e-41, 1.4013e-45, 0.0000e+00, -1.4677e+24,
- 4.5915e-41]], requires_grad=True)
- """
- if not issubclass(module_cls, torch.nn.Module):
- raise RuntimeError(f'Expected a Module; got {module_cls}')
- if 'device' not in inspect.signature(module_cls).parameters:
- raise RuntimeError('Module must support a \'device\' arg to skip initialization')
- final_device = kwargs.pop('device', 'cpu')
- kwargs['device'] = 'meta'
- return module_cls(*args, **kwargs).to_empty(device=final_device)
|
