patch bug in IAccountant len functions. Currently returns # of unique noise_multiplier values vs intended # of steps

opacus/accountants/accountant.py

@@ -51,13 +51,6 @@ def get_epsilon(self, delta: float, *args, **kwargs) -> float:
         """
         pass
 
-    @abc.abstractmethod
-    def __len__(self) -> int:
-        """
-        Number of optimization steps taken so far
-        """
-        pass
-
     @classmethod
     @abc.abstractmethod
     def mechanism(cls) -> str:
@@ -66,6 +59,13 @@ def mechanism(cls) -> str:
         """
         pass
 
+
+    def __len__(self) -> int:
+        """
+        Number of optimization steps taken so far
+        """
+        return sum(num_steps for _, _, num_steps in self.history)
+
     def get_optimizer_hook_fn(
         self, sample_rate: float
     ) -> Callable[[DPOptimizer], None]:

opacus/accountants/gdp.py

@@ -67,9 +67,6 @@ def get_epsilon(self, delta: float, poisson: bool = True, **kwargs) -> float:
             delta=delta,
         )
 
-    def __len__(self):
-        return len(self.history)
-
     @classmethod
     def mechanism(cls) -> str:
         return "gdp"

opacus/accountants/prv.py

@@ -157,6 +157,3 @@ def _get_domain(
     @classmethod
     def mechanism(cls) -> str:
         return "prv"
-
-    def __len__(self):
-        return len(self.history)

opacus/accountants/rdp.py

@@ -84,9 +84,6 @@ def get_epsilon(
         eps, _ = self.get_privacy_spent(delta=delta, alphas=alphas)
         return eps
 
-    def __len__(self):
-        return len(self.history)
-
     @classmethod
     def mechanism(cls) -> str:
         return "rdp"

opacus/tests/accountants_test.py

@@ -129,6 +129,30 @@ def test_prv_accountant(self) -> None:
         epsilon = accountant.get_epsilon(delta=1e-5)
         self.assertAlmostEqual(epsilon, 6.777395712150674)
 
+    def test_len_counts_optimization_steps(self) -> None:
+        noise_multiplier = 1.5
+        sample_rate = 0.04
+        steps = 50
+
+        for accountant in (RDPAccountant(), GaussianAccountant(), PRVAccountant()):
+            for _ in range(steps):
+                accountant.step(
+                    noise_multiplier=noise_multiplier, sample_rate=sample_rate
+                )
+            self.assertEqual(len(accountant.history), 1)
+            self.assertEqual(len(accountant), steps)
+
+        for accountant in (RDPAccountant(), PRVAccountant()):
+            for _ in range(steps):
+                accountant.step(
+                    noise_multiplier=noise_multiplier, sample_rate=sample_rate
+                )
+            accountant.step(
+                noise_multiplier=noise_multiplier * 2, sample_rate=sample_rate
+            )
+            self.assertEqual(len(accountant.history), 2)
+            self.assertEqual(len(accountant), steps + 1)
+
     def test_get_noise_multiplier_rdp_epochs(self) -> None:
         delta = 1e-5
         sample_rate = 0.04