94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
|
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)
|