A class to represent the memory state of a model layer. Args: layer_name (str): The name of the layer. layer_mem (int): The memory usage of the layer in bytes.
| 13 | |
| 14 | |
| 15 | class SimpleMemState: |
| 16 | """ |
| 17 | A class to represent the memory state of a model layer. |
| 18 | |
| 19 | Args: |
| 20 | layer_name (str): The name of the layer. |
| 21 | layer_mem (int): The memory usage of the layer in bytes. |
| 22 | """ |
| 23 | |
| 24 | def __init__(self, layer_name: str, layer_mem: int = 0) -> None: |
| 25 | self.layer_name = layer_name |
| 26 | |
| 27 | # Memory status of the current model layer. |
| 28 | self._layer_mem: int = layer_mem |
| 29 | # Total memory status of the model and sub-models, initialized with layer memory. |
| 30 | self._total_mem: int = self._layer_mem |
| 31 | # SimpleMemState of sub-models. |
| 32 | self.sub_model_stats = OrderedDict() |
| 33 | |
| 34 | @property |
| 35 | def layer_mem(self) -> int: |
| 36 | """ |
| 37 | Get the memory usage of the layer. |
| 38 | |
| 39 | Returns: |
| 40 | int: The memory usage of the layer in bytes. |
| 41 | """ |
| 42 | return self._layer_mem |
| 43 | |
| 44 | @layer_mem.setter |
| 45 | def layer_mem(self, new_layer_mem: int) -> None: |
| 46 | """ |
| 47 | Set the memory usage of the layer. |
| 48 | |
| 49 | Args: |
| 50 | new_layer_mem (int): The new memory usage of the layer in bytes. |
| 51 | """ |
| 52 | diff = new_layer_mem - self._layer_mem |
| 53 | self._layer_mem = new_layer_mem |
| 54 | self._total_mem += diff |
| 55 | |
| 56 | @property |
| 57 | def total_mem(self) -> int: |
| 58 | """ |
| 59 | Get the total memory usage of the model and sub-models. |
| 60 | |
| 61 | Returns: |
| 62 | int: The total memory usage in bytes. |
| 63 | """ |
| 64 | return self._total_mem |
| 65 | |
| 66 | def add(self, layer_name: str, layer_mem: int = 0, flush: bool = True) -> None: |
| 67 | """ |
| 68 | Add a layer to the memory state. |
| 69 | |
| 70 | Args: |
| 71 | layer_name (str): The name of the layer. |
| 72 | layer_mem (int, optional): The memory usage of the layer in bytes. Defaults to 0. |
no outgoing calls
no test coverage detected