Use Python properties to cleanly apply rules to class fields

Summary: As titled. There are strong rules on most of the fields in influence algorithms. Make use of Python properties to make the structure cleaner and enforce rules even when client code tries to change the field.

Reviewed By: yucu

Differential Revision: D54663278

fbshipit-source-id: 273a82e4707c1b0a7afb8637459b4e34bf381d89

Author: cicichen01

captum/influence/_core/tracincp_fast_rand_proj.py

@@ -186,9 +186,6 @@ def __init__(
 
         # TODO: restore prior state
         self.final_fc_layer = final_fc_layer
-        if isinstance(self.final_fc_layer, str):
-            self.final_fc_layer = _get_module_from_name(model, self.final_fc_layer)
-        assert isinstance(self.final_fc_layer, Module)
         for param in self.final_fc_layer.parameters():
             param.requires_grad = True
 
@@ -203,6 +200,24 @@ def __init__(
             else _check_loss_fn(self, test_loss_fn, "test_loss_fn")
         )
 
+    @property
+    def final_fc_layer(self) -> Module:
+        return self._final_fc_layer
+
+    @final_fc_layer.setter
+    def final_fc_layer(self, layer: Union[Module, str]):
+        if isinstance(layer, str):
+            try:
+                self._final_fc_layer = _get_module_from_name(self.model, layer)
+                if not isinstance(self._final_fc_layer, Module):
+                    raise Exception("No module found for final_fc_layer")
+            except Exception as ex:
+                raise ValueError(
+                    f'Invalid final_fc_layer str: "{layer}" provided!'
+                ) from ex
+        else:
+            self._final_fc_layer = layer
+
     @log_usage()
     def influence(  # type: ignore[override]
         self,

tests/influence/_core/test_tracin_validation.py

@@ -16,8 +16,7 @@
 
 class TestTracinValidator(BaseTest):
 
-    param_list = []
-    for reduction, constr in [
+    param_list = [
         (
             "none",
             DataInfluenceConstructor(TracInCP, name="TracInCP"),
@@ -29,8 +28,7 @@ class TestTracinValidator(BaseTest):
                 name="TracInCpFast",
             ),
         ),
-    ]:
-        param_list.append((reduction, constr))
+    ]
 
     @parameterized.expand(
         param_list,
@@ -64,3 +62,24 @@ def test_tracin_require_inputs_dataset(
             )
             with self.assertRaisesRegex(AssertionError, "required."):
                 tracin.influence(None, k=None)
+
+    def test_tracincp_fast_rand_proj_inputs(self) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            (
+                net,
+                train_dataset,
+                test_samples,
+                test_labels,
+            ) = get_random_model_and_data(tmpdir, unpack_inputs=False)
+
+            with self.assertRaisesRegex(
+                ValueError, 'Invalid final_fc_layer str: "invalid_layer" provided!'
+            ):
+                TracInCPFast(
+                    net,
+                    "invalid_layer",
+                    train_dataset,
+                    tmpdir,
+                    loss_fn=nn.MSELoss(),
+                    batch_size=1,
+                )