Core

Base class for all ATOMMIC models

class atommic.core.classes.modelPT.ModelPT(*args: Any, **kwargs: Any)

Bases: LightningModule, Model

Interface for Pytorch-lightning based atommic models.

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

Register model artifacts with this function. These artifacts (files) will be included inside .atommic file when model.save_to(“model.atommic”) 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 "atommic_file:unique_artifact_name":

    .atommic will be untarred to a temporary folder location and an actual existing path will be
    returned

else:

    an error will be raised.

If “src” is local existing path, then it will be returned in absolute path form. elif “src” starts with “atommic_file:unique_artifact_name” .atommic 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

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

• src (str) – Path to artifact.

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

Returns

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

• life.

has_artifacts() → bool

Returns True if model has artifacts registered.

has_native_or_submodules_artifacts() → bool

Returns True if it has artifacts or any of the submodules have artifacts.

has_atommic_submodules() → bool

Returns True if it has any registered atommic submodules.

register_atommic_submodule(name: str, config_field: str, model: ModelPT) → None

Adds a atommic model as a submodule.

Submodule can be accessed via the name attribute on the parent atommic model this submodule was registered on (self).

In the saving process, the whole parent model (self) is held as a solid model with artifacts from the child submodule, the submodule config will be saved to the config_field of the parent model.

This method is necessary to create a nested model, e.g.

class ParentModel(ModelPT):
    def __init__(self, cfg, trainer=None):
        super().__init__(cfg=cfg, trainer=trainer)

        # annotate type for autocompletion and type checking (optional)
        self.child_model: Optional[ChildModel] = None
        if cfg.get("child_model") is not None:
            self.register_atommic_submodule(
                name="child_model",
                config_field="child_model",
                model=ChildModel(self.cfg.child_model, trainer=trainer),
            )
        # ... other code

Parameters

• name (str) – Name of the submodule. This name will be used to access the submodule from the parent model.

• config_field (str) – Name of the config field where the submodule config will be saved.

• model (ModelPT) – The submodule model.

named_atommic_modules(prefix_name: str = '', prefix_config: str = '') → Iterator[Tuple[str, str, ModelPT]]

Returns an iterator over all atommic submodules recursively, yielding tuples of (attribute path, path in config, submodule), starting from the core module.

Parameters

• prefix_name (str) – Prefix for the name path.

• prefix_config (str) – Prefix for the path in config.

Return type

Iterator over (attribute path, path in config, submodule), starting from (prefix, self).

save_to(save_path: str)

Saves model instance (weights and configuration) into .atommic file. You can use “restore_from” method to fully restore instance from .atommic file. .atommic 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.ckpt - model checkpoint

Parameters

saved. (Path to .atommic file where model instance should be) –

classmethod restore_from(restore_path: str, override_config_path: Optional[Union[omegaconf.OmegaConf, str]] = None, map_location: Optional[torch.device] = None, strict: bool = True, return_config: bool = False, save_restore_connector: Optional[SaveRestoreConnector] = None, trainer: Optional[pytorch_lightning.Trainer] = None)

Restores model instance (weights and configuration) from .atommic file.

Parameters

• restore_path (str) – Path to .atommic file from which model should be instantiated override_config_path: path to a yaml config that will override the internal config file or an OmegaConf/DictConfig object representing the model config.

• override_config_path (Optional[Union[OmegaConf, str]]) – path to a yaml config that will override the internal config file or an OmegaConf/DictConfig object representing the model config.

• map_location (Optional[torch.device]) – Optional torch.device() to map the instantiated model to a device. Default is None, it will select a GPU if available, falling back to CPU otherwise.

• strict (bool) – Passed to load_state_dict. Default is True.

• return_config (bool) – If set to true, will return just the underlying config of the restored model as an OmegaConf/DictConfig object without instantiating the model.

• save_restore_connector (SaveRestoreConnector) – Can be overridden to add custom save and restore logic.

• trainer (Optional[Trainer]) – Optional, a pytorch lightning Trainer object that will be forwarded to the instantiated model’s constructor.

Example

model = atommic.collections.reconstruction.nn.CIRIM.restore_from('CIRIM.atommic')
assert isinstance(model, atommic.collections.reconstruction.nn.CIRIM)

Return type

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

classmethod load_from_checkpoint(checkpoint_path: str, *args, map_location: Optional[Union[Dict[str, str], str, torch.device, int, Callable]] = None, hparams_file: Optional[str] = None, strict: bool = True, **kwargs)

Loads ModelPT from checkpoint, with some maintenance of restoration. For documentation, please refer to LightningModule.load_from_checkpoint() documentation.

Parameters

• checkpoint_path (str) – Path to checkpoint. Will be passed to LightningModule.load_from_checkpoint().

• *args – Will be passed to LightningModule.load_from_checkpoint().

• map_location (Optional[Union[Dict[str, str], str, torch.device, int, Callable]]) – Will be passed to LightningModule.load_from_checkpoint().

• hparams_file (Optional[str]) – Will be passed to LightningModule.load_from_checkpoint().

• strict (bool) – Will be passed to LightningModule.load_from_checkpoint().

abstract setup_training_data(train_data_config: Union[omegaconf.DictConfig, Dict])

Setups data loader to be used in training.

abstract setup_validation_data(val_data_config: Union[omegaconf.DictConfig, Dict])

Setups data loader to be used in validation.

setup_test_data(test_data_config: Union[omegaconf.DictConfig, Dict])

(Optionally) Setups data loader to be used in test.

setup_multiple_validation_data(val_data_config: Union[omegaconf.DictConfig, Dict])

(Optionally) Setups data loader to be used in validation.

setup_multiple_test_data(test_data_config: Union[omegaconf.DictConfig, Dict])

(Optionally) Setups data loader to be used in test, with support for multiple data loaders.

setup_optimization(optim_config: Optional[Union[omegaconf.DictConfig, Dict]] = None)

Prepares an optimizer from a string name and its optional config parameters.

Parameters

optim_config (Optional[Union[DictConfig, Dict]]) –

A dictionary containing the following keys:

• lrfloat

  Mandatory key for learning rate. Will raise ValueError if not provided.

• optimizerstr

  String name pointing to one of the available optimizers in the registry. Default is Adam.

• opt_argslist

  Optional list of strings, in the format “arg_name=arg_value”. The list of “arg_value” will be parsed and a dictionary of optimizer kwargs will be built and supplied to instantiate the optimizer.

Return type

An instance of an optimizer.

setup_optimizer_param_groups()

Used to create param groups for the optimizer. As an example, this can be used to specify per-layer learning rates:

optim.SGD([
            {'params': model.base.parameters()},
            {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

See https://pytorch.org/docs/stable/optim.html for more information. By default, ModelPT will use self.parameters(). Override this method to add custom param groups.

configure_optimizers()

Configure optimizers and schedulers for training.

setup(stage: Optional[str] = None)

Called at the beginning of fit, validate, test, or predict. This is called on every process when using DDP.

Parameters

stage (str) – fit, validate, test or predict

train_dataloader()

Return the training dataloader.

val_dataloader()

Return the validation dataloader.

test_dataloader()

Return the test dataloader.

on_validation_epoch_end() → Optional[Dict[str, Dict[str, torch.Tensor]]]

Default DataLoader for Validation set which automatically supports multiple data loaders via multi_validation_epoch_end. If multi dataset support is not required, override this method entirely in base class. In such a case, there is no need to implement multi_validation_epoch_end either.

Note

If more than one data loader exists, and they all provide val_loss, only the val_loss of the first data loader will be used by default. This default can be changed by passing the special key val_dl_idx: int inside the validation_ds config.

Returns

• A dictionary containing the union of all items from individual data_loaders, along with merged logs from all

• data loaders.

on_test_epoch_end() → Optional[Dict[str, Dict[str, torch.Tensor]]]

Default DataLoader for Test set which automatically supports multiple data loaders via multi_test_epoch_end. If multi dataset support is not required, override this method entirely in base class. In such a case, there is no need to implement multi_test_epoch_end either.

Note

If more than one data loader exists, and they all provide test_loss, only the test_loss of the first data loader will be used by default. This default can be changed by passing the special key _test_dl_idx: int inside the test_ds config.

Returns

A dictionary containing the union of all items from individual data_loaders, along with merged logs from all data loaders.

Return type

Dict[str, Dict[str, torch.Tensor]]

static multi_validation_epoch_end(outputs: Optional[Union[object, List[Dict[str, torch.Tensor]]]], dataloader_idx: int = 0) → None

Adds support for multiple validation datasets. Should be overridden by subclass, to obtain appropriate logs for each of the dataloaders.

Parameters

• outputs (Union[List[Dict[str, torch.Tensor]], List[List[Dict[str, torch.Tensor]]]]) – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.

• dataloader_idx (int) – The index of the dataloader for which the outputs are provided.

Returns

• A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be

• pre-pended by the dataloader prefix.

static multi_test_epoch_end(outputs: Union[object, List[Dict[str, torch.Tensor]]], dataloader_idx: int = 0) → None

Adds support for multiple test datasets. Should be overridden by subclass, to obtain appropriate logs for each of the dataloaders.

Parameters

• outputs (Union[List[Dict[str, torch.Tensor]], List[List[Dict[str, torch.Tensor]]]]) – Same as that provided by LightningModule.on_test_epoch_end() for a single dataloader.

• dataloader_idx (int) – The index of the dataloader for which the outputs are provided.

Returns

• A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be

• pre-pended by the dataloader prefix.

get_validation_dataloader_prefix(dataloader_idx: int = 0) → str

Get the name of one or more data loaders, which will be prepended to all logs.

get_test_dataloader_prefix(dataloader_idx: int = 0) → str

Get the name of one or more data loaders, which will be prepended to all logs.

load_part_of_state_dict(state_dict, include, exclude, load_from_string=None)

Load part of the state dict.

maybe_init_from_pretrained_checkpoint(cfg: omegaconf.OmegaConf, map_location: str = 'cpu')

Initializes a given model with the parameters obtained via specific config arguments. The state dict of the provided model will be updated with strict=False setting to prevent requirement of exact model parameters matching.

Initializations

init_from_atommic_modelstr

Str path to a .atommic model, which will be instantiated in order to extract the state dict.

init_from_pretrained_modelstr

Str name of a pretrained model checkpoint (obtained via cloud). The model will be downloaded (or a cached copy will be used), instantiated and then its state dict will be extracted.

init_from_ptl_ckptstr

Str name of a Pytorch Lightning checkpoint file. It will be loaded and the state dict will extract.

param cfg

The config used to instantiate the model. It needs only contain one of the above keys.

type cfg

OmegaConf

param map_location

str or torch.device() which represents where the intermediate state dict (from the pretrained model or checkpoint) will be loaded.

type map_location

str or torch.device()

classmethod extract_state_dict_from(restore_path: str, save_dir: str, split_by_module: bool = False, save_restore_connector: Optional[SaveRestoreConnector] = None)

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

Parameters

• restore_path (str) – Path to .atommic file from which state dict(s) should be extracted.

• save_dir (str) – Directory in which the saved state dict(s) should be stored.

• split_by_module (bool, optional) – Bool flag, which determines whether the output checkpoint should be for the entire Model, or the individual module’s that comprise the Model. Default is False.

• save_restore_connector (SaveRestoreConnector, optional) – Can be overridden to add custom save and restore logic. Default is None.

Example

To convert the .atommic tarfile into a single Model level PyTorch checkpoint

state_dict = atommic.collections.asr.models.EncDecCTCModel.extract_state_dict_from('asr.atommic',             './asr_ckpts')

To restore a model from a Model level checkpoint

model = atommic.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 .atommic tarfile into multiple Module level PyTorch checkpoints

state_dict = atommic.collections.asr.models.EncDecCTCModel.extract_state_dict_from('asr.atommic',             './asr_ckpts', split_by_module=True)

To restore a module from a Module level checkpoint

model = atommic.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 .atommic checkpoint.

prepare_test(trainer: pytorch_lightning.Trainer) → bool

Helper method to check whether the model can safely be tested on a dataset after training (or loading a checkpoint).

trainer = Trainer()
if model.prepare_test(trainer):
    trainer.test(model)

Return type

Bool which declares the model safe to test. Provides warnings if it has to return False to guide the user.

set_trainer(trainer: pytorch_lightning.Trainer)

Set an instance of Trainer object.

set_world_size(trainer: pytorch_lightning.Trainer)

Determines the world size from the PyTorch Lightning Trainer and then updates AppState.

summarize(max_depth: int = 1) → pytorch_lightning.utilities.model_summary.ModelSummary

Summarize this LightningModule.

property num_weights

Utility property that returns the total number of parameters of the Model.

property cfg

Property that holds the finalized internal config of the model.

Note

Changes to this config are not reflected in the state of the model. Please create a new model using an updated config to properly update the model.

classmethod update_save_restore_connector(save_restore_connector)

Update the save_restore_connector of the model.

on_train_start()

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-start We use it here to copy the relevant config for dynamic freezing.

on_train_batch_start(batch: Any, batch_idx: int, unused: int = 0)

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-batch-start We use it here to enable nsys profiling and dynamic freezing.

on_train_batch_end(outputs, batch: Any, batch_idx: int, unused: int = 0) → None

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-batch-end We use it here to enable nsys profiling.

on_train_end()

PyTorch Lightning hook: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#on-train-end We use it here to clean up the dynamic freezing config.

cuda(device=None)

PTL is overriding this method and changing the pytorch behavior of a module. The PTL LightingModule override will move the module to device 0 if device is None.

See the PTL method here: Lightning-AI/lightning device_dtype_mixin.py#L113

Here we are overriding this to maintain the default Pytorch nn.module behavior: pytorch/pytorch

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters

device (torch.device, optional) – The desired device of the parameters and buffers in this module.

Base Mixin classes

class atommic.core.classes.common.Typing

Bases: ABC

An interface which endows module with neural types.

property input_types: Optional[Dict[str, NeuralType]]

Define these to enable input neural type checks.

property output_types: Optional[Dict[str, NeuralType]]

Define these to enable output neural type checks.

_validate_input_types(input_types=None, ignore_collections=False, **kwargs)

This function does a few things.

1. It ensures that len(self.input_types <non-optional>) <= len(kwargs) <= len(self.input_types).

2. For each (keyword name, keyword value) passed as input to the wrapped function:

   • Check if the keyword name exists in the list of valid self.input_types names.

   • Check if keyword value has the neural_type property.

     • If it does, then perform a comparative check and assert that neural types are compatible (SAME

       or GREATER).

   • Check if keyword value is a container type (list or tuple). If yes, then perform the elementwise

     test of neural type above on each element of the nested structure, recursively.

Parameters

• input_types (class) – Either the input_types defined at class level, or the local function overridden type definition.

• ignore_collections (bool) –

  For backward compatibility, container support can be disabled explicitly using this flag. When set to True,

  all nesting is ignored and nest-depth checks are skipped.

• kwargs (Dict[str, Any]) – Dictionary of argument_name:argument_value pairs passed to the wrapped function upon call.

_attach_and_validate_output_types(out_objects, ignore_collections=False, output_types=None)

This function does a few things.

1. It ensures that len(out_object) == len(self.output_types).

2. If the output is a tensor (or list/tuple of list/tuple … of tensors), it

   attaches a neural_type to it. For objects without the neural_type attribute, such as python objects (dictionaries and lists, primitive data types, structs), no neural_type is attached. Note: tensor.neural_type is only checked during _validate_input_types which is called prior to forward().

Parameters

• output_types (class) – Either the output_types defined at class level, or the local function overridden type definition.

• ignore_collections (bool) –

  For backward compatibility, container support can be disabled explicitly using this flag. When set to True,

  all nesting is ignored and nest-depth checks are skipped.

• out_objects (Dict[str, Any]) – The outputs of the wrapped function.

__check_neural_type(obj, metadata, depth: int, name: Optional[str] = None)

Checks if the object is of the correct type, and attaches the correct NeuralType.

Parameters

• obj (object) – Any python object that can be assigned to a value.

• metadata (object) – TypecheckMetadata object.

• depth (int) – Current depth of the recursion.

• name (str) – Optional name used of the source object, when an error is raised.

__attach_neural_type(obj, metadata, depth: int, name: Optional[str] = None)

Attach NeuralType to the object.

Parameters

• obj (object) – Any python object that can be assigned to a value.

• metadata (object) – TypecheckMetadata object.

• depth (int) – Current depth of the recursion.

• name (str) – Optional name used of the source object, when an error is raised.

---

class atommic.core.classes.common.Serialization

Bases: ABC

Base class for serialization.

classmethod from_config_dict(config: omegaconf.DictConfig, trainer: Optional[pytorch_lightning.Trainer] = None)

Instantiates object using DictConfig-based configuration

to_config_dict() → omegaconf.DictConfig

Returns object’s configuration to config dictionary

---

class atommic.core.classes.common.FileIO

Bases: ABC

Base class for file IO.

save_to(save_path: str)

Standardized method to save a tarfile containing the checkpoint, config, and any additional artifacts. Implemented via atommic.core.connectors.save_restore_connector.SaveRestoreConnector.save_to().

Parameters

save_path (str) – Path to save the checkpoint to.

classmethod restore_from(restore_path: str, override_config_path: Optional[str] = None, map_location: Optional[torch.device] = None, strict: bool = True, return_config: bool = False, trainer: Optional[pytorch_lightning.Trainer] = None, save_restore_connector: Optional[SaveRestoreConnector] = None)

Restores model instance (weights and configuration) from a .atommic file.

Parameters

• restore_path (str) – Path to .atommic file from which model should be instantiated.

• override_config_path (str, optional) – Path to .yaml file containing the configuration to override the one in the .atommic file.

• map_location (torch.device, optional) – Device to map the instantiated model to. Default is None, it will select a GPU if available, falling back to CPU otherwise.

• strict (bool, optional) – Passed to load_state_dict. Default is True.

• return_config (bool, optional) – If True, returns the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model.

• trainer (Trainer, optional) – If provided, will be used to instantiate the model.

• save_restore_connector (SaveRestoreConnector, optional) – An optional SaveRestoreConnector object that defines the implementation of the restore_from() method.

classmethod from_config_file(path2yaml_file: str)

Instantiates an instance of atommic Model from YAML config file. Weights will be initialized randomly.

Parameters

path2yaml_file (str) – Path to yaml file with model configuration.

Return type

atommic Model instance.

to_config_file(path2yaml_file: str)

Saves current instance’s configuration to YAML config file. Weights will not be saved.

Parameters

path2yaml_file (str) – Path to yaml file with model configuration.

Base Connector classes

class atommic.core.connectors.save_restore_connector.SaveRestoreConnector

Bases: object

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

save_to(model: atommic.ModelPT, save_path: str)

Saves model instance (weights and configuration) into .atommic file.

You can use “restore_from” method to fully restore instance from .atommic file.

.atommic 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_weights.ckpt - model checkpoint

Parameters

• model (ModelPT) – ModelPT object to be saved.

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

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

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

Parameters

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

• restore_path (str) – Path to .atommic file from which model should be instantiated

• override_config_path (Optional[Union[OmegaConf, str]]) – Path to a yaml config that will override the internal config file or an OmegaConf/DictConfig object representing the model config.

• map_location (Optional[torch.device]) – Optional torch.device() to map the instantiated model to a device. Default is None, it will select a GPU if available, falling back to CPU otherwise.

• strict (bool) – Passed to load_state_dict. When set to False, the model will be able to load a checkpoint that has more parameters than the model itself. Default is True.

• return_config (bool) – If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model.

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

Example

` model = atommic.collections.asr.models.EncDecCTCModel.restore_from('asr.atommic') assert isinstance(model, atommic.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.

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

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

Parameters

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

• restore_path (str) – Path to .atommic file from which model should be instantiated.

• override_config_path (str or OmegaConf/DictConfig object, optional) – Path to a yaml config that will override the internal config file or an OmegaConf/DictConfig object representing the model config.

• map_location (torch.device, optional) – Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise.

• strict (bool, optional) – Passed to load_state_dict. Default is True.

• return_config (bool, optional) – If set to true, will return just the underlying config of the restored model as an OmegaConf/DictConfig object without instantiating the model.

• trainer (Trainer, optional) – 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).

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

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

Parameters

• restore_path (str) – Path to .atommic file from which state dict(s) should be extracted.

• save_dir (str) – Directory in which the saved state dict(s) should be stored.

• split_by_module (bool, optional) – Bool flag, which determines whether the output checkpoint should be for the entire Model, or the individual module’s that comprise the Model. Default is False.

Example

To convert the .atommic tarfile into a single Model level PyTorch checkpoint :: state_dict = atommic.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.atommic’, ‘./asr_ckpts’) To restore a model from a Model level checkpoint :: model = atommic.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 .atommic tarfile into multiple Module level PyTorch checkpoints :: state_dict = atommic.collections.asr.models.EncDecCTCModel.extract_state_dict_from(‘asr.atommic’, ‘./asr_ckpts’, split_by_module=True). To restore a module from a Module level checkpoint :: model = atommic.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 .atommic checkpoint.

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 .atommic file when model.save_to(“mymodel.atommic”) 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 “atommic_file:unique_artifact_name”: .atommic will be untarred to a temporary folder location and an actual existing path will be returned else an error will be raised.

If "src" is local existing path:
    then it will be returned in absolute path form
elif "src" starts with "atommic_file:unique_artifact_name":
    .atommic 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) – ModelPT object to register artifact for.

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

• src (str) – Path to artifact.

• verify_src_exists (bool, optional) – If set to False, then the artifact is optional and register_artifact will return None even if src is not found. Default is 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

property model_config_yaml: str

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

property model_weights_ckpt: str

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

property model_extracted_dir: Optional[str]

Neural Type checking

class atommic.core.classes.common.typecheck(input_types: Optional[Union[TypeState, Dict[str, NeuralType]]] = TypeState.UNINITIALIZED, output_types: Optional[Union[TypeState, Dict[str, NeuralType]]] = TypeState.UNINITIALIZED, ignore_collections: bool = False)

Bases: object

A decorator which performs input-output neural type checks, and attaches neural types to the output of the function that it wraps. Requires that the class inherit from atommic.core.Typing in order to perform type checking, and will raise an error if that is not the case.

# Usage (Class level type support) .. code-block:: python

@typecheck() def fn(self, arg1, arg2, …):

# Usage (Function level type support) .. code-block:: python

@typecheck(input_types=…, output_types=…) def fn(self, arg1, arg2, …):

Points to be noted:

1. The brackets () in @typecheck() are necessary. You will encounter a TypeError: __init__() takes 1 positional argument but X were given without those brackets.

2. The function can take any number of positional arguments during definition. When you call this function, all arguments must be passed using kwargs only.

__call__(enabled=None, adapter=None, proxy=<class 'FunctionWrapper'>)

Wrapper method that can be used on any function of a class that implements Typing. By default, it will utilize the input_types and output_types properties of the class inheriting Typing. Local function level overrides can be provided by supplying dictionaries as arguments to the decorator.

class TypeState(value)

Bases: Enum

Placeholder to denote the default value of type information provided. If the constructor of this decorator is used to override the class level type definition, this enum value indicate that types will be overridden.

UNINITIALIZED = 0

static set_typecheck_enabled(enabled: bool = True)

Set the global typecheck flag.

static disable_checks()

Temporarily disable type checks.

Neural Type classes

class atommic.core.neural_types.neural_type.NeuralType(axes: Optional[Tuple] = None, elements_type: ElementType = VoidType, optional=False)

Bases: object

This is the main class which would represent neural type concept. It is used to represent the types of inputs and outputs.

compare(second) → NeuralTypeComparisonResult

Performs neural type comparison of self with second. When you chain two modules’ inputs/outputs via __call__ method, this comparison will be called to ensure neural type compatibility.

compare_and_raise_error(parent_type_name, port_name, second_object)

Method compares definition of one type with another and raises an error if not compatible.

---

class atommic.core.neural_types.axes.AxisType(kind: AxisKindAbstract, size: Optional[int] = None, is_list=False)

Bases: object

This class represents axis semantics and (optionally) it’s dimensionality.

---

class atommic.core.neural_types.elements.ElementType

Bases: ABC

Abstract class defining semantics of the tensor elements. We are relying on Python for inheritance checking

property type_parameters: Dict

Override this property to parametrize your type. For example, you can specify ‘storage’ type such as float, int, bool with ‘dtype’ keyword. Another example, is if you want to represent a signal with a particular property (say, sample frequency), then you can put sample_freq->value in there. When two types are compared their type_parameters must match.”

property fields: Optional[Tuple]

This should be used to logically represent tuples/structures. For example, if you want to represent a bounding box (x, y, width, height) you can put a tuple with names (‘x’, y’, ‘w’, ‘h’) in here. Under the hood this should be converted to the last tensor dimension of fixed size = len(fields). When two types are compared their fields must match.

compare(second) → NeuralTypeComparisonResult

Override this method to provide a comparison between two types.

---

class atommic.core.neural_types.comparison.NeuralTypeComparisonResult(value)

Bases: Enum

The result of comparing two neural type objects for compatibility. When comparing A.compare_to(B).

SAME = 0

LESS = 1

GREATER = 2

DIM_INCOMPATIBLE = 3

TRANSPOSE_SAME = 4

INCOMPATIBLE = 6

SAME_TYPE_INCOMPATIBLE_PARAMS = 7

Experiment manager

class atommic.utils.exp_manager.exp_manager(trainer: pytorch_lightning.Trainer, cfg: Optional[Union[omegaconf.DictConfig, Dict]] = None)

Bases:

exp_manager is a helper function used to manage folders for experiments. It follows the pytorch lightning paradigm of exp_dir/model_or_experiment_name/version. If the lightning trainer has a logger, exp_manager will get exp_dir, name, and version from the logger. Otherwise, it will use the exp_dir and name arguments to create the logging directory. exp_manager also allows for explicit folder creation via explicit_log_dir.

The version can be a datetime string or an integer. Date time version can be disabled if use_datetime_version is set to False. It optionally creates TensorBoardLogger, WandBLogger, ModelCheckpoint objects from pytorch lightning. It copies sys.argv, and git information if available to the logging directory. It creates a log file for each process to log their output into.

exp_manager additionally has a resume feature (resume_if_exists) which can be used to continuing training from the constructed log_dir. When you need to continue the training repeatedly (like on a cluster which you need multiple consecutive jobs), you need to avoid creating the version folders. Therefore, from v1.0.0, when resume_if_exists is set to True, creating the version folders is ignored.

Parameters

• trainer (pytorch_lightning.Trainer) – The lightning trainer object.

• cfg (DictConfig or Dict, optional) –

  Can have the following keys:

  • explicit_log_dirstr

    Can be used to override exp_dir/name/version folder creation. Defaults to None, which will use

    exp_dir, name, and version to construct the logging directory.

  • exp_dirstr

    The base directory to create the logging directory. Defaults to None, which logs to

    ./atommic_experiments.

  • namestr

    The name of the experiment. Defaults to None which turns into “default” via name = name or

    ”default”.

  • versionstr

    The version of the experiment. Defaults to None which uses either a datetime string or lightning’s

    TensorboardLogger system of using version_{int}.

  • use_datetime_versionbool

    Whether to use a datetime string for version. Default is True.

  • resume_if_existsbool

    Whether this experiment is resuming from a previous run. If True, it sets

    trainer._checkpoint_connector._ckpt_path so that the trainer should auto-resume. exp_manager will move files under log_dir to log_dir/run_{int}. Default is False. When resume_if_exists is True, we would not create version folders to make it easier to find the log folder for next runs.

  • resume_past_endbool

    exp_manager errors out if resume_if_exists is True and a checkpoint matching ‘'*end.ckpt indicating a

    previous training run fully completed. This behaviour can be disabled, in which case the ‘'*end.ckpt will be loaded by setting resume_past_end to True. Default is False.

  • resume_ignore_no_checkpointbool

    exp_manager errors out if resume_if_exists is True and no checkpoint could be found. This behaviour

    can be disabled, in which case exp_manager will print a message and continue without restoring, by setting resume_ignore_no_checkpoint to True. Default is False.

  • resume_from_checkpointstr

    Can be used to specify a path to a specific checkpoint file to load from. This will override any

    checkpoint found when resume_if_exists is True. Default is None.

  • create_tensorboard_loggerbool

    Whether to create a tensorboard logger and attach it to the pytorch lightning trainer.

    Default is True.

  • summary_writer_kwargsdict

    A dictionary of kwargs that can be passed to lightning’s TensorboardLogger class. Note that log_dir is

    passed by exp_manager and cannot exist in this dict. Default is None.

  • create_wandb_loggerbool

    Whether to create a Weights and Biases logger and attach it to the pytorch lightning trainer.

    Default is False.

  • wandb_logger_kwargsdict

    A dictionary of kwargs that can be passed to lightning’s WandBLogger class. Note that name and project

    are required parameters if create_wandb_logger is True. Default is None..

  • create_checkpoint_callbackbool

    Whether to create a ModelCheckpoint callback and attach it to the pytorch lightning trainer. The

    ModelCheckpoint saves the top 3 models with the best “val_loss”, the most recent checkpoint under ‘'*last.ckpt, and the final checkpoint after training completes under ‘'*end.ckpt. Default is True.

  • create_early_stopping_callbackbool

    Whether to create an EarlyStopping callback and attach it to the pytorch lightning trainer. The

    EarlyStopping callback stops training if the “val_loss” does not improve for 3 epochs. Default is True.

  • files_to_copylist

    A list of files to copy to the experiment logging directory. Defaults to None which copies no files.

  • log_local_rank_0_onlybool

    Whether to only create log files for local rank 0. Default is False. Set this to True if you are

    using DDP with many GPUs and do not want many log files in your exp dir.

  • log_global_rank_0_onlybool

    Whether to only create log files for global rank 0. Defaults to False. Set this to True if you are

    using DDP with many GPUs and do not want many log files in your exp dir.

  • max_timestr

    The maximum wall clock time per run. This is intended to be used on clusters where you want a

    checkpoint to be saved after this specified time and be able to resume from that checkpoint. Default is None.

  • seconds_to_sleepfloat

    Seconds to sleep non rank 0 processes for. Used to give enough time for rank 0 to initialize.

Returns

log_dir – The final logging directory where logging files are saved. Usually the concatenation of exp_dir, name, and version.

Return type

Path

class atommic.utils.exp_manager.ExpManagerConfig(explicit_log_dir: Optional[str] = None, exp_dir: Optional[str] = None, name: Optional[str] = None, version: Optional[str] = None, use_datetime_version: Optional[bool] = True, resume_if_exists: Optional[bool] = False, resume_past_end: Optional[bool] = False, resume_ignore_no_checkpoint: Optional[bool] = False, resume_from_checkpoint: Optional[str] = None, create_tensorboard_logger: Optional[bool] = True, summary_writer_kwargs: Optional[Dict[Any, Any]] = None, create_wandb_logger: Optional[bool] = False, wandb_logger_kwargs: Optional[Dict[Any, Any]] = None, create_checkpoint_callback: Optional[bool] = True, checkpoint_callback_params: Optional[CallbackParams] = CallbackParams(filepath=None, dirpath=None, filename=None, monitor='val_loss', verbose=True, save_last=True, save_top_k=3, save_weights_only=False, mode='min', auto_insert_metric_name=True, every_n_epochs=1, every_n_train_steps=None, train_time_interval=None, prefix=None, postfix='.atommic', save_best_model=False, always_save_atommic=False, save_atommic_on_train_end=True, model_parallel_size=None, save_on_train_epoch_end=False), create_early_stopping_callback: Optional[bool] = False, early_stopping_callback_params: Optional[EarlyStoppingParams] = EarlyStoppingParams(monitor='val_loss', mode='min', min_delta=0.001, patience=10, verbose=True, strict=True, check_finite=True, stopping_threshold=None, divergence_threshold=None, check_on_train_epoch_end=None, log_rank_zero_only=False), create_preemption_callback: Optional[bool] = True, files_to_copy: Optional[List[str]] = None, log_step_timing: Optional[bool] = True, step_timing_kwargs: Optional[StepTimingParams] = StepTimingParams(reduction='mean', sync_cuda=False, buffer_size=1), log_local_rank_0_only: Optional[bool] = False, log_global_rank_0_only: Optional[bool] = False, disable_validation_on_resume: Optional[bool] = True, ema: Optional[EMAParams] = EMAParams(enable=False, decay=0.999, cpu_offload=False, validate_original_weights=False, every_n_steps=1), max_time_per_run: Optional[str] = None, seconds_to_sleep: float = 5)

Bases: object

Configuration for the experiment manager.

explicit_log_dir: Optional[str] = None

exp_dir: Optional[str] = None

name: Optional[str] = None

version: Optional[str] = None

use_datetime_version: Optional[bool] = True

resume_if_exists: Optional[bool] = False

resume_past_end: Optional[bool] = False

resume_ignore_no_checkpoint: Optional[bool] = False

resume_from_checkpoint: Optional[str] = None

create_tensorboard_logger: Optional[bool] = True

summary_writer_kwargs: Optional[Dict[Any, Any]] = None

create_wandb_logger: Optional[bool] = False

wandb_logger_kwargs: Optional[Dict[Any, Any]] = None

create_checkpoint_callback: Optional[bool] = True

checkpoint_callback_params: Optional[CallbackParams] = CallbackParams(filepath=None, dirpath=None, filename=None, monitor='val_loss', verbose=True, save_last=True, save_top_k=3, save_weights_only=False, mode='min', auto_insert_metric_name=True, every_n_epochs=1, every_n_train_steps=None, train_time_interval=None, prefix=None, postfix='.atommic', save_best_model=False, always_save_atommic=False, save_atommic_on_train_end=True, model_parallel_size=None, save_on_train_epoch_end=False)

create_early_stopping_callback: Optional[bool] = False

early_stopping_callback_params: Optional[EarlyStoppingParams] = EarlyStoppingParams(monitor='val_loss', mode='min', min_delta=0.001, patience=10, verbose=True, strict=True, check_finite=True, stopping_threshold=None, divergence_threshold=None, check_on_train_epoch_end=None, log_rank_zero_only=False)

create_preemption_callback: Optional[bool] = True

files_to_copy: Optional[List[str]] = None

log_step_timing: Optional[bool] = True

step_timing_kwargs: Optional[StepTimingParams] = StepTimingParams(reduction='mean', sync_cuda=False, buffer_size=1)

log_local_rank_0_only: Optional[bool] = False

log_global_rank_0_only: Optional[bool] = False

disable_validation_on_resume: Optional[bool] = True

ema: Optional[EMAParams] = EMAParams(enable=False, decay=0.999, cpu_offload=False, validate_original_weights=False, every_n_steps=1)

max_time_per_run: Optional[str] = None

seconds_to_sleep: float = 5

Exportable

class atommic.core.classes.export.Exportable

Bases: ABC

This Interface should be implemented by particular classes derived from atommic.core.ModelPT. It gives these entities ability to be exported for deployment to formats such as ONNX.

property input_module

Implement this method to return the input module

property output_module

Implement this method to return the output module.

export(output: str, input_example=None, verbose=False, do_constant_folding=True, onnx_opset_version=None, check_trace: Union[bool, List[torch.Tensor]] = False, dynamic_axes=None, check_tolerance=0.01, export_modules_as_functions: bool = False, keep_initializers_as_inputs=None)

Export the module to a file.

Parameters

• output (str) – The output file path.

• input_example (dict) – A dictionary of input names and values.

• verbose (bool) – If True, print out the export process.

• do_constant_folding (bool) – If True, do constant folding.

• onnx_opset_version (int) – The ONNX opset version to use.

• check_trace (bool or list of torch.Tensor) – If True, check the trace of the exported model.

• dynamic_axes (dict) – A dictionary of input names and dynamic axes.

• check_tolerance (float) – The tolerance for the check_trace.

• export_modules_as_functions (bool) – If True, export modules as functions.

• keep_initializers_as_inputs (bool) – If True, keep initializers as inputs.

property disabled_deployment_input_names

Implement this method to return a set of input names disabled for export

property disabled_deployment_output_names

Implement this method to return a set of output names disabled for export

property supported_export_formats

Implement this method to return a set of export formats supported. Default is all types.

property input_names

Implement this method to return a list of input names

property output_names

Override this method to return a set of output names disabled for export

property input_types_for_export

Implement this method to return a list of input types

property output_types_for_export

Implement this method to return a list of output types

get_export_subnet(subnet=None)

Returns Exportable subnet model/module to export

static list_export_subnets()

Returns default set of subnet names exported for this model. First goes the one receiving input (input_example).

get_export_config()

Returns export_config dictionary.

set_export_config(args)

Sets/updates export_config dictionary.