feat(ctx): define data attribution for AnalogContext with read-only protection

src/aihwkit/nn/conversion.py

@@ -138,6 +138,7 @@ def convert_to_analog(
 
         verbose: Increase verbosity. Will print converted layers.
 
+
     Returns:
         Module where all the digital layers are replaced with analog
         mapped layers.

src/aihwkit/nn/modules/base.py

@@ -7,10 +7,12 @@
 """Base class for adding functionality to analog layers."""
 from typing import Any, List, Optional, Tuple, NamedTuple, Union, Generator, Callable, TYPE_CHECKING
 from collections import OrderedDict
+import warnings
 
 from torch import Tensor
 from torch.nn import Parameter
 from torch import device as torch_device
+from torch import Size as torch_size
 
 from aihwkit.exceptions import ModuleError
 from aihwkit.simulator.tiles.module import TileModule
@@ -244,6 +246,30 @@ class type when setting ``load_rpu_config`` to
             ``load_rpu_config=False``.
 
         """
+        keys_to_delete = []
+        for name, param in list(state_dict.items()):
+            if name.endswith("analog_ctx") and param.size() == torch_size([]):
+                keys_to_delete.append(name)
+
+        # For the checkpoint saved by aihwkit before version, 0.9.2, the parameters with
+        #    class `AnalogContext` are saved as empty size tensors since `AnalogContext`,
+        #    which is a derived class of `torch.nn.Parameter`,
+        #    uses an empty Parameter tensor to store the context.
+        if len(keys_to_delete) > 0:
+            strict = False
+            for key in keys_to_delete:
+                del state_dict[key]
+            warnings.warn(
+                "Some parameters in the loaded checkpoint has empty size"
+                "(param.size() == torch.Size([]))."
+                "It could happens because of the loaded checkpoint"
+                "is generated by an older version of aihwkit."
+                "The parameter is skipped for compatibility reasons."
+                "The loading mode is set to non-strict."
+                "It is recommended to re-save the checkpoint with the latest version of aihwkit."
+                "Related parameters are: {}".format(keys_to_delete)
+            )
+
         for analog_tile in self.analog_tiles():
             analog_tile.set_load_rpu_config_state(load_rpu_config, strict_rpu_config_check)
         return super().load_state_dict(state_dict, strict)  # type: ignore

src/aihwkit/nn/modules/linear.py

@@ -82,6 +82,10 @@ def reset_parameters(self) -> None:
             self.weight, self.bias = self.get_weights()  # type: ignore
             super().reset_parameters()
             self.set_weights(self.weight, self.bias)  # type: ignore
+            # AnalogLinear doesn't support access weight and bias directly, so delete them
+            del self.weight, self.bias
+            # delete them manually is necessary since asigning `bias` (a bool) is forbidden
+            #   by torch if self.bias is already a tensor
             self.weight, self.bias = None, bias
 
     def forward(self, x_input: Tensor) -> Tensor: