mridc.core.connectors package

Submodules

mridc.core.connectors.save_restore_connector module

class mridc.core.connectors.save_restore_connector.SaveRestoreConnector

Bases: object

This class is used to save and restore the model state.

extract_state_dict_from(restore_path: str, save_dir: str, split_by_module: bool = False)

Extract the state dict(s) from a provided .mridc tarfile and save it to a directory.

Parameters

• restore_path (path to .mridc file from which state dict(s) should be extracted) –

• save_dir (directory in which the saved state dict(s) should be stored) –

• split_by_module (bool flag, which determines whether the output checkpoint should be for the entire Model, or) –

• Model. (the individual module's that comprise the) –

Example

To convert the .mridc tarfile into a single Model level PyTorch checkpoint :: state_dict = mridc.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.mridc’, ‘./asr_ckpts’) To restore a model from a Model level checkpoint :: model = mridc.collections.asr.models.EncDecCTCModel(cfg) # or any other method of restoration model.load_state_dict(torch.load(“./asr_ckpts/model_weights.ckpt”)) To convert the .mridc tarfile into multiple Module level PyTorch checkpoints :: state_dict = mridc.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.mridc’, ‘./asr_ckpts’, split_by_module=True). To restore a module from a Module level checkpoint :: model = mridc.collections.asr.models.EncDecCTCModel(cfg) # or any other method of restoration # load the individual components model.preprocessor.load_state_dict(torch.load(“./asr_ckpts/preprocessor.ckpt”)) model.encoder.load_state_dict(torch.load(“./asr_ckpts/encoder.ckpt”)) model.decoder.load_state_dict(torch.load(“./asr_ckpts/decoder.ckpt”))

Return type

The state dict that was loaded from the original .mridc checkpoint.

load_config_and_state_dict(calling_cls, restore_path: str, override_config_path: Optional[Union[OmegaConf, str]] = None, map_location: Optional[device] = None, strict: bool = True, return_config: bool = False, trainer: Optional[Trainer] = None)

Restores model instance (weights and configuration) into .mridc file

Parameters

• calling_cls (Class of the model to be restored.) –

• restore_path (path to .mridc file from which model should be instantiated) –

• override_config_path (path to a yaml config that will override the internal config file or an) –

• config. (OmegaConf/DictConfig object representing the model) –

• map_location (Optional torch.device() to map the instantiated model to a device. By default (None), it will) –

• available (select a GPU if) –

• otherwise. (falling back to CPU) –

• strict (Passed to load_state_dict. By default, True.) –

• return_config (If set to true, will return just the underlying config of the restored model as an OmegaConf) –

• model. (DictConfig object without instantiating the) –

• trainer (Optional trainer object to be used for model parallelism.) –

Example

` model = mridc.collections.asr.models.EncDecCTCModel.restore_from('asr.mridc') assert isinstance(model, mridc.collections.asr.models.EncDecCTCModel) `

Return type

An instance of type cls or its underlying config (if return_config is set).

static load_instance_with_state_dict(instance, state_dict, strict)

Loads the state dict into the instance.

property model_config_yaml: str

This property is used to get the path to the model config yaml file.

property model_extracted_dir: Optional[str]

property model_weights_ckpt: str

This property is used to get the path to the model weights ckpt file.

static register_artifact(model, config_path: str, src: str, verify_src_exists: bool = True)

Register model artifacts with this function. These artifacts (files) will be included inside .mridc file when model.save_to(“mymodel.mridc”) is called.

How it works: 1. It always returns existing absolute path which can be used during Model constructor call. EXCEPTION: src is None or “” in which case nothing will be done and src will be returned 2. It will add (config_path, model_utils.ArtifactItem()) pair to self.artifacts. If “src” is local existing path, then it will be returned in absolute path form. elif “src” starts with “mridc_file:unique_artifact_name”: .mridc will be untarred to a temporary folder location and an actual existing path will be returned else an error will be raised.

WARNING: use .register_artifact calls in your models’ constructors. The returned path is not guaranteed to exist after you have exited your model’s constructor.

Parameters

• model (ModelPT object to register artifact for.) –

• config_path (Artifact key. Usually corresponds to the model config.) –

• src (Path to artifact.) –

• verify_src_exists (If set to False, then the artifact is optional and register_artifact will return None) – even if src is not found. Defaults to True.

Returns

life.

Return type

If src is not None or empty it always returns absolute path which is guaranteed to exist during model instance

restore_from(calling_cls, restore_path: str, override_config_path: Optional[Union[OmegaConf, str]] = None, map_location: Optional[device] = None, strict: bool = True, return_config: bool = False, trainer: Optional[Trainer] = None)

Restores model instance (weights and configuration) into .mridc file

Parameters

• calling_cls (The class of the model to be restored.) –

• restore_path (path to .mridc file from which model should be instantiated) –

• override_config_path (path to a yaml config that will override the internal config file or an) –

• config. (OmegaConf/DictConfig object representing the model) –

• map_location (Optional torch.device() to map the instantiated model to a device. By default (None), it will) –

• available (select a GPU if) –

• otherwise. (falling back to CPU) –

• strict (Passed to load_state_dict. By default, True.) –

• return_config (If set to true, will return just the underlying config of the restored model as an) –

• model. (OmegaConf/DictConfig object without instantiating the) –

• trainer (Optional trainer object to be used for restoring the model.) –

Return type

An instance of type cls or its underlying config (if return_config is set).

save_to(model, save_path: str)

Saves model instance (weights and configuration) into .mridc file. You can use “restore_from” method to fully restore instance from .mridc file. .mridc file is an archive (tar.gz) with the following: - model_config.yaml - model configuration in .yaml format. You can deserialize this into cfg argument for model’s constructor - model_wights.chpt - model checkpoint

Parameters

• model (ModelPT object to be saved.) –

• save_path (Path to .mridc file where model instance should be saved) –

Module contents