| from contextlib import contextmanager
|
| import torch
|
| import torch.nn as nn
|
|
|
| @contextmanager
|
| def init_empty_weights(include_buffers: bool=False):
|
| """Meta initialization context manager.
|
|
|
| A context manager under which models are initialized with all parameters
|
| on the meta device, therefore creating an empty model. Useful when just
|
| initializing the model would blow the available RAM.
|
|
|
| Args:
|
| include_buffers (`bool`, *optional*, defaults to `False`): Whether or
|
| not to also put all buffers on the meta device while initializing.
|
|
|
| Example:
|
| ```python
|
| import torch.nn as nn
|
|
|
| # Initialize a model with 100 billions parameters in no time and without using any RAM.
|
| with init_empty_weights():
|
| tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
|
| ```
|
|
|
| <Tip warning={true}>
|
|
|
| Any model created under this context manager has no weights. As such you can't do something like
|
| `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
|
|
|
| </Tip>
|
| """
|
| with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
|
| yield f
|
|
|
| @contextmanager
|
| def init_on_device(device: torch.device, include_buffers: bool=False):
|
| """Device initialization context manager.
|
|
|
| A context manager under which models are initialized with all parameters
|
| on the specified device.
|
|
|
| Args:
|
| device (`torch.device`): Device to initialize all parameters on.
|
| include_buffers (`bool`, *optional*, defaults to `False`): Whether or
|
| not to also put all buffers on the meta device while initializing.
|
|
|
| Example:
|
| ```python
|
| import torch.nn as nn
|
|
|
| with init_on_device(device=torch.device("cuda")):
|
| tst = nn.Liner(100, 100) # on `cuda` device
|
| ```
|
| """
|
| old_register_parameter = nn.Module.register_parameter
|
| if include_buffers:
|
| old_register_buffer = nn.Module.register_buffer
|
|
|
| def register_empty_parameter(module, name, param):
|
| old_register_parameter(module, name, param)
|
| if param is not None:
|
| param_cls = type(module._parameters[name])
|
| kwargs = module._parameters[name].__dict__
|
| module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
|
|
|
| def register_empty_buffer(module, name, buffer):
|
| old_register_buffer(module, name, buffer)
|
| if buffer is not None:
|
| module._buffers[name] = module._buffers[name].to(device)
|
| if include_buffers:
|
| tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
|
| else:
|
| tensor_constructors_to_patch = {}
|
|
|
| def patch_tensor_constructor(fn):
|
|
|
| def wrapper(*args, **kwargs):
|
| kwargs['device'] = device
|
| return fn(*args, **kwargs)
|
| return wrapper
|
| try:
|
| nn.Module.register_parameter = register_empty_parameter
|
| if include_buffers:
|
| nn.Module.register_buffer = register_empty_buffer
|
| for torch_function_name in tensor_constructors_to_patch.keys():
|
| setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
|
| yield
|
| finally:
|
| nn.Module.register_parameter = old_register_parameter
|
| if include_buffers:
|
| nn.Module.register_buffer = old_register_buffer
|
| for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
|
| setattr(torch, torch_function_name, old_torch_function) |
