Updated dev dependencies and tests

.github/workflows/python-package.yml

@@ -31,6 +31,7 @@ jobs:
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
         # Installing pytorch on cpu and torchvision for running tests 
         python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+        python -m pip install matplotlib
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

pyproject.toml

@@ -42,3 +42,8 @@ docs = [
     "mkdocstrings[python]",
     "mkdocs-material",
 ]
+dev = [
+    "matplotlib",
+    "pytest",
+    "flake8"
+]

tests/test_ChainedBSScheduler.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import ChainedBSScheduler, ConstantBS, ExponentialBS
@@ -67,7 +68,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("ChainedBSScheduler.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ChainedBSScheduler.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -86,7 +88,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("ChainedScheduler.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ChainedScheduler.png")
         plt.close()
 
 

tests/test_ConstantBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import ConstantBS
@@ -63,7 +64,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("ConstantBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ConstantBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -79,7 +81,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("ConstantLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ConstantLR.png")
         plt.close()
 
 

tests/test_CosineAnnealingBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import CosineAnnealingBS
@@ -69,7 +70,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("CosineAnnealingBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CosineAnnealingBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -85,7 +87,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("CosineAnnealingLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CosineAnnealingLR.png")
         plt.close()
 
 

tests/test_CosineAnnealingBSWithWarmRestarts.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import CosineAnnealingBSWithWarmRestarts
@@ -66,7 +67,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("CosineAnnealingBSWithWarmRestarts.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CosineAnnealingBSWithWarmRestarts.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -83,7 +85,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("CosineAnnealingWarmRestarts.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CosineAnnealingWarmRestarts.png")
         plt.close()
 
 

tests/test_CyclicBS.py

@@ -1,3 +1,4 @@
+import os
 import random
 import unittest
 
@@ -160,7 +161,8 @@ def test_graphic_triangular2(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("CyclicBS-triangular2.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CyclicBS-triangular2.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -177,7 +179,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("CyclicLR-triangular2.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CyclicLR-triangular2.png")
         plt.close()
 
     def test_graphic_exp_range(self):
@@ -197,7 +200,8 @@ def test_graphic_exp_range(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("CyclicBS-exp_range.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CyclicBS-exp_range.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -214,7 +218,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("CyclicLR-exp_range.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/CyclicLR-exp_range.png")
         plt.close()
 
 

tests/test_ExponentialBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import ExponentialBS
@@ -57,7 +58,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("ExponentialBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ExponentialBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -73,7 +75,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("ExponentialLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ExponentialLR.png")
         plt.close()
 
 

tests/test_IncreaseBSOnPlateau.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import IncreaseBSOnPlateau
@@ -77,7 +78,8 @@ def test_graphic(self):
         scheduler = IncreaseBSOnPlateau(dataloader, mode='min', threshold_mode='rel', max_batch_size=max_batch_size)
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, metrics)
         plt.plot(batch_sizes)
-        plt.savefig("IncreaseBSOnPlateau.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/IncreaseBSOnPlateau.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -95,7 +97,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step(**d)
         plt.plot(learning_rates)
-        plt.savefig("ReduceLROnPlateau.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/ReduceLROnPlateau.png")
         plt.close()
 
 

tests/test_LinearBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 import torch
@@ -94,7 +95,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("LinearBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/LinearBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -111,7 +113,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("LinearLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/LinearLR.png")
         plt.close()
 
 

tests/test_MultiStepBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import MultiStepBS
@@ -75,7 +76,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("MultiStepBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/MultiStepBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -91,7 +93,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("MultiStepLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/MultiStepLR.png")
         plt.close()
 
 

tests/test_OneCycleBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import OneCycleBS
@@ -94,7 +95,8 @@ def test_graphic(self):
         # TODO: not equivalent, OneCycleBS can step more than total_steps
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("OneCycleBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/OneCycleBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -111,7 +113,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("OneCycleLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/OneCycleLR.png")
         plt.close()
 
 

tests/test_PolynomialBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import PolynomialBS
@@ -63,7 +64,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("PolynomialBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/PolynomialBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -79,7 +81,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("PolynomialLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/PolynomialLR.png")
         plt.close()
 
 

tests/test_StepBS.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 
 from bs_scheduler import StepBS
@@ -76,7 +77,8 @@ def test_graphic(self):
 
         batch_sizes = get_batch_sizes_across_epochs(dataloader, scheduler, n_epochs)
         plt.plot(batch_sizes)
-        plt.savefig("StepBS.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/StepBS.png")
         plt.close()
 
         model = torch.nn.Linear(10, 10)
@@ -92,7 +94,8 @@ def get_lr(optimizer):
             learning_rates.append(get_lr(optimizer))
             scheduler.step()
         plt.plot(learning_rates)
-        plt.savefig("StepLR.png")
+        os.makedirs("images", exist_ok=True)
+        plt.savefig("images/StepLR.png")
         plt.close()
 
 

tests/test_utils.py

@@ -4,9 +4,9 @@
 import unittest
 
 import torch
+from torch.utils.data import DataLoader
 from torchvision import datasets
 from torchvision.transforms import ToTensor
-from torch.utils.data import DataLoader
 
 from bs_scheduler import BSScheduler