Fixing ExponentialBS rounding errors for small gamma (#13)

* Previous behavior: when the exponential factor `gamma` used by `ExponentialBS` is close to 1, the batch size would remain constant. * Fixed this by maintaining a float value for the batch size as an internal state. * Updated tests. * Added new plots. * Bumping version.
ancestor-mithril · Jun 5, 2024 · 11be404 · 11be404
1 parent 2bed4ae
commit 11be404
Show file tree

Hide file tree

Showing 29 changed files with 14 additions and 8 deletions.
diff --git a/bs_scheduler/batch_size_schedulers.py b/bs_scheduler/batch_size_schedulers.py
@@ -665,6 +665,7 @@ def __init__(self, dataloader: DataLoader, gamma: float, batch_size_manager: Uni
         assert gamma > 0.0
         # Gamma is expected to be greater than 1.0 for batch size growth. It can be lower than 1.0 for batch size decay.
         self.gamma: float = gamma
+        self.float_bs: Union[float, None] = None
         super().__init__(dataloader, batch_size_manager, max_batch_size, min_batch_size, verbose)
 
     def get_new_bs(self) -> int:
@@ -674,7 +675,12 @@ def get_new_bs(self) -> int:
         if self.last_epoch == 0:
             return self.batch_size
 
-        return rint(self.batch_size * self.gamma)
+        if self.float_bs is None or rint(self.float_bs) != self.batch_size:
+            # Using rint instead of int because otherwise we will increas the BS faster
+            self.float_bs = self.batch_size
+
+        self.float_bs *= self.gamma
+        return rint(self.float_bs)
 
 
 class SequentialBS(BSScheduler):

diff --git a/docs/img/plots/ChainedBSScheduler.png b/docs/img/plots/ChainedBSScheduler.png
diff --git a/docs/img/plots/ChainedScheduler.png b/docs/img/plots/ChainedScheduler.png
diff --git a/docs/img/plots/ConstantBS.png b/docs/img/plots/ConstantBS.png
diff --git a/docs/img/plots/ConstantLR.png b/docs/img/plots/ConstantLR.png
diff --git a/docs/img/plots/CosineAnnealingBS.png b/docs/img/plots/CosineAnnealingBS.png
diff --git a/docs/img/plots/CosineAnnealingBSWithWarmRestarts.png b/docs/img/plots/CosineAnnealingBSWithWarmRestarts.png
diff --git a/docs/img/plots/CosineAnnealingLR.png b/docs/img/plots/CosineAnnealingLR.png
diff --git a/docs/img/plots/CosineAnnealingWarmRestarts.png b/docs/img/plots/CosineAnnealingWarmRestarts.png
diff --git a/docs/img/plots/CyclicBS-exp_range.png b/docs/img/plots/CyclicBS-exp_range.png
diff --git a/docs/img/plots/CyclicBS-triangular2.png b/docs/img/plots/CyclicBS-triangular2.png
diff --git a/docs/img/plots/CyclicLR-exp_range.png b/docs/img/plots/CyclicLR-exp_range.png
diff --git a/docs/img/plots/CyclicLR-triangular2.png b/docs/img/plots/CyclicLR-triangular2.png
diff --git a/docs/img/plots/ExponentialBS.png b/docs/img/plots/ExponentialBS.png
diff --git a/docs/img/plots/ExponentialLR.png b/docs/img/plots/ExponentialLR.png
diff --git a/docs/img/plots/IncreaseBSOnPlateau.png b/docs/img/plots/IncreaseBSOnPlateau.png
diff --git a/docs/img/plots/LinearBS.png b/docs/img/plots/LinearBS.png
diff --git a/docs/img/plots/LinearLR.png b/docs/img/plots/LinearLR.png
diff --git a/docs/img/plots/MultiStepBS.png b/docs/img/plots/MultiStepBS.png
diff --git a/docs/img/plots/MultiStepLR.png b/docs/img/plots/MultiStepLR.png
diff --git a/docs/img/plots/OneCycleBS.png b/docs/img/plots/OneCycleBS.png
diff --git a/docs/img/plots/OneCycleLR.png b/docs/img/plots/OneCycleLR.png
diff --git a/docs/img/plots/ReduceLROnPlateau.png b/docs/img/plots/ReduceLROnPlateau.png
diff --git a/docs/img/plots/StepBS.png b/docs/img/plots/StepBS.png
diff --git a/docs/img/plots/StepLR.png b/docs/img/plots/StepLR.png
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "bs_scheduler"
-version = "0.4.1"
+version = "0.4.2"
 requires-python = ">=3.9"
 description = "A PyTorch Dataloader compatible batch size scheduler library."
 readme = "README.md"

diff --git a/tests/test_ChainedBSScheduler.py b/tests/test_ChainedBSScheduler.py
@@ -24,7 +24,7 @@ def test_dataloader_lengths(self):
         n_epochs = 10
 
         epoch_lengths = simulate_n_epochs(dataloader, scheduler, n_epochs)
-        expected_batch_sizes = [100, 110, 121, 133, 14, 15, 16, 18, 20, 22]
+        expected_batch_sizes = [100, 110, 121, 133, 14, 16, 17, 19, 21, 23]
         expected_lengths = self.compute_epoch_lengths(expected_batch_sizes, len(self.dataset), drop_last=False)
         self.assertEqual(epoch_lengths, expected_lengths)
 
@@ -37,7 +37,7 @@ def test_dataloader_batch_size(self):
         n_epochs = 10
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
-        expected_batch_sizes = [100, 110, 121, 133, 14, 15, 16, 18, 20, 22]
+        expected_batch_sizes = [100, 110, 121, 133, 14, 16, 17, 19, 21, 23]
 
         self.assertEqual(batch_sizes, expected_batch_sizes)
 

diff --git a/tests/test_ExponentialBS.py b/tests/test_ExponentialBS.py
@@ -53,8 +53,8 @@ def test_graphic(self):
 
         base_batch_size = 10
         dataloader = create_dataloader(self.dataset, batch_size=base_batch_size)
-        scheduler = ExponentialBS(dataloader, gamma=2, max_batch_size=100, verbose=False)
-        n_epochs = 10
+        scheduler = ExponentialBS(dataloader, gamma=1.05, max_batch_size=500, verbose=False)
+        n_epochs = 50
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
@@ -64,7 +64,7 @@ def test_graphic(self):
 
         model = torch.nn.Linear(10, 10)
         optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
-        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5)
+        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
         learning_rates = []
 
         def get_lr(optimizer):

diff --git a/tests/test_SequentialBS.py b/tests/test_SequentialBS.py
@@ -24,7 +24,7 @@ def test_dataloader_lengths(self):
         n_epochs = 10
 
         epoch_lengths = simulate_n_epochs(dataloader, scheduler, n_epochs)
-        expected_batch_sizes = [100] * 4 + [10, 11, 12, 13, 14, 15]
+        expected_batch_sizes = [100] * 4 + [10, 11, 12, 13, 15, 16]
         expected_lengths = self.compute_epoch_lengths(expected_batch_sizes, len(self.dataset), drop_last=False)
         self.assertEqual(epoch_lengths, expected_lengths)