| | |
| | 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) |
| |
|
