65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
import functools
|
|
from dora.log import fatal
|
|
import warnings
|
|
import torch
|
|
from pathlib import Path
|
|
import inspect
|
|
|
|
def _check_diffq():
|
|
try:
|
|
import diffq
|
|
except ImportError:
|
|
fatal('Trying to use DiffQ, but diffq is not installed.\n'
|
|
'On Windows run: python.exe -m pip install diffq \n'
|
|
'On Linux/Mac, run: python3 -m pip install diffq')
|
|
|
|
def capture_init(init):
|
|
@functools.wraps(init)
|
|
def __init__(self, *args, **kwargs):
|
|
self._init_args_kwargs = (args, kwargs)
|
|
init(self, *args, **kwargs)
|
|
|
|
return __init__
|
|
|
|
|
|
def load_model(path_or_package, strict=False):
|
|
if isinstance(path_or_package, dict):
|
|
package = path_or_package
|
|
elif isinstance(path_or_package, (str, Path)):
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
path = path_or_package
|
|
package = torch.load(path, 'cpu')
|
|
else:
|
|
raise ValueError(f"Invalid type for {path_or_package}.")
|
|
|
|
klass = package["klass"]
|
|
args = package["args"]
|
|
kwargs = package["kwargs"]
|
|
|
|
if strict:
|
|
model = klass(*args, **kwargs)
|
|
else:
|
|
sig = inspect.signature(klass)
|
|
for key in list(kwargs):
|
|
if key not in sig.parameters:
|
|
warnings.warn("Dropping inexistant parameter " + key)
|
|
del kwargs[key]
|
|
model = klass(*args, **kwargs)
|
|
|
|
state = package["state"]
|
|
|
|
set_state(model, state)
|
|
return model
|
|
|
|
def set_state(model, state, quantizer=None):
|
|
if state.get('__quantized'):
|
|
if quantizer is not None:
|
|
quantizer.restore_quantized_state(model, state['quantized'])
|
|
else:
|
|
_check_diffq()
|
|
from diffq import restore_quantized_state
|
|
restore_quantized_state(model, state)
|
|
else:
|
|
model.load_state_dict(state)
|
|
return state |