Skip to content

cassetta.io

Overview

Thos module contains routines for data input/output.

It also defines mixins and utilities for making load/unloading of modules and checkpoints easier than in pure PyTorch.

MODULE DESCRIPTION
modules

Advanced loading/unloading of modules & models

cassetta.io.modules

LoadableModule

Bases: LoadableMixin, Module

A Loadable variant of nn.Module

LoadableSequential 

LoadableSequential(*args)

Bases: LoadableMixin, Sequential

A Loadable variant of nn.Sequential

Source code in cassetta/io/modules.py 
64
65
66
def __init__(self, *args) -> None:
    super().__init__(*args)
    _validate_loadable_modules(self)

LoadableModuleList 

LoadableModuleList(modules=None)

Bases: LoadableMixin, ModuleList

A Loadable variant of nn.ModuleList

Source code in cassetta/io/modules.py 
97
98
99
def __init__(self, modules: Optional[Iterable[nn.Module]] = None) -> None:
    super().__init__(modules)
    _validate_loadable_modules(self)

LoadableModuleDict 

LoadableModuleDict(modules=None)

Bases: LoadableMixin, ModuleDict

A Loadable variant of nn.ModuleDict

Source code in cassetta/io/modules.py 
135
136
137
138
139
140
def __init__(self, modules: Optional[_ValidInput] = None) -> None:
    super().__init__(modules)
    # Ensure all modules are loadable
    for key, module in self.items():
        if not isinstance(module, LoadableMixin):
            raise TypeError(f"Module '{key}' must be Loadable")