[docs]defskip_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) """ifnotissubclass(module_cls,torch.nn.Module):raiseRuntimeError(f"Expected a Module; got {module_cls}")if"device"notininspect.signature(module_cls).parameters:raiseRuntimeError("Module must support a 'device' arg to skip initialization")final_device=kwargs.pop("device","cpu")kwargs["device"]="meta"returnmodule_cls(*args,**kwargs).to_empty(device=final_device)
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.