|
|
|
|
|
|
|
|
""" |
|
|
@File : states.py |
|
|
@Time : 2023/8/8 下午7:01 |
|
|
@Author : waytan |
|
|
@Contact : waytan@tencent.com |
|
|
@License : (C)Copyright 2023, Tencent |
|
|
@Desc : Utilities to save and load models. |
|
|
""" |
|
|
import functools |
|
|
import inspect |
|
|
import warnings |
|
|
from pathlib import Path |
|
|
from fractions import Fraction |
|
|
|
|
|
import torch |
|
|
|
|
|
|
|
|
def load_state_dict(net, pth_path): |
|
|
kwargs = {'sources': ['drums', 'bass', 'other', 'vocal'], 'audio_channels': 2, 'samplerate': 44100, |
|
|
'segment': Fraction(39, 5), 'channels': 48, 'channels_time': None, 'growth': 2, 'nfft': 4096, |
|
|
'wiener_iters': 0, 'end_iters': 0, 'wiener_residual': False, 'cac': True, 'depth': 4, 'rewrite': True, |
|
|
'multi_freqs': [], 'multi_freqs_depth': 3, 'freq_emb': 0.2, 'emb_scale': 10, 'emb_smooth': True, |
|
|
'kernel_size': 8, 'stride': 4, 'time_stride': 2, 'context': 1, 'context_enc': 0, 'norm_starts': 4, |
|
|
'norm_groups': 4, 'dconv_mode': 3, 'dconv_depth': 2, 'dconv_comp': 8, 'dconv_init': 0.001, |
|
|
'bottom_channels': 512, 't_layers': 5, 't_hidden_scale': 4.0, 't_heads': 8, 't_dropout': 0.02, |
|
|
't_layer_scale': True, 't_gelu': True, 't_emb': 'sin', 't_max_positions': 10000, 't_max_period': 10000.0, |
|
|
't_weight_pos_embed': 1.0, 't_cape_mean_normalize': True, 't_cape_augment': True, |
|
|
't_cape_glob_loc_scale': [5000.0, 1.0, 1.4], 't_sin_random_shift': 0, 't_norm_in': True, |
|
|
't_norm_in_group': False, 't_group_norm': False, 't_norm_first': True, 't_norm_out': True, |
|
|
't_weight_decay': 0.0, 't_lr': None, 't_sparse_self_attn': False, 't_sparse_cross_attn': False, |
|
|
't_mask_type': 'diag', 't_mask_random_seed': 42, 't_sparse_attn_window': 400, 't_global_window': 100, |
|
|
't_sparsity': 0.95, 't_auto_sparsity': False, 't_cross_first': False, 'rescale': 0.1} |
|
|
model = net(**kwargs) |
|
|
state_dict = torch.load(pth_path) |
|
|
model.load_state_dict(state_dict) |
|
|
return model |
|
|
|
|
|
|
|
|
def load_model(path_or_package, strict=False): |
|
|
"""Load a model from the given serialized model, either given as a dict (already loaded) |
|
|
or a path to a file on disk.""" |
|
|
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 get_state(model, quantizer, half=False): |
|
|
"""Get the state from a model, potentially with quantization applied. |
|
|
If `half` is True, model are stored as half precision, which shouldn't impact performance |
|
|
but half the state size.""" |
|
|
if quantizer is None: |
|
|
dtype = torch.half if half else None |
|
|
state = {k: p.data.to(device='cpu', dtype=dtype) for k, p in model.state_dict().items()} |
|
|
else: |
|
|
state = quantizer.get_quantized_state() |
|
|
state['__quantized'] = True |
|
|
return state |
|
|
|
|
|
|
|
|
def set_state(model, state, quantizer=None): |
|
|
"""Set the state on a given model.""" |
|
|
if state.get('__quantized'): |
|
|
quantizer.restore_quantized_state(model, state['quantized']) |
|
|
else: |
|
|
model.load_state_dict(state) |
|
|
return state |
|
|
|
|
|
|
|
|
def capture_init(init): |
|
|
@functools.wraps(init) |
|
|
def __init__(self, *args, **kwargs): |
|
|
self._init_args_kwargs = (args, kwargs) |
|
|
init(self, *args, **kwargs) |
|
|
|
|
|
return __init__ |
|
|
|
