Added lr scheduler nodes and cifar10 workflows.

2025-03-12 14:57:44 -07:00
parent 635c97c037
commit acb5d9f820
20 changed files with 770 additions and 6 deletions
@@ -0,0 +1,77 @@
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.ptn_conv_model import PtnConvModel
+from pytorch_wrapper.pto_adam import PtoAdam
+from pytorch_wrapper.ptv_dataset import PtvDataset
+from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
+from pytorch_wrapper.pt_data_loader import PtDataLoader
+from pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr
+from pytorch_wrapper.pt_save_model import PtSaveModel
+from pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep  # Updated import
+
+
+class TestPtTrainClassificationConvModelLr(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        compose_node = PtvTransformsToTensor()
+        transform = compose_node.f()[0]
+        dataset_node = PtvDataset()
+        dataset = dataset_node.f("FashionMNIST",
+                    True, # download
+                    "", # root
+                    '{"train": True}',
+                    transform=transform)[0]
+        data_loader_node = PtDataLoader()
+        self.train_loader = data_loader_node.f(
+            dataset=dataset,
+            batch_size=32,
+            shuffle=True,
+            parameters='{"num_workers":1}'
+        )[0]
+
+        self.model_node = PtnConvModel()
+        self.model = self.model_node.f(
+                 input_dim="(1, 28, 28)",
+                 penultimate_dim=0,
+                 output_dim=10,
+                 channel_list="[32,64,128,256]",
+                 kernel_size_list="[3,3,3,3]",
+                 padding_list="[1,1,1,1]",
+                 downsample_list="[True,True,True,True]"
+        )[0]
+
+        self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0]
+        
+        self.scheduler = PtoLrSchedulerStep().f(self.optimizer,
+                                                3,  # steps
+                                                0.1  # gamma
+                                               )[0]
+
+        self.trainer = PtTrainClassificationModelLr()
+        self.save_model = PtSaveModel()
+
+    def test_1(self):
+        epochs = 4
+        trained_model = self.trainer.f(
+                        self.model,
+                        self.train_loader,
+                        self.optimizer,
+                        self.scheduler,
+                        epochs, # epochs
+                        use_gpu=True,
+                        early_stopping=False,
+                        early_stopping_rounds=1,
+                        output_best_val_model=False)[0]
+        self.save_model.f(trained_model, f"fashion_mnist_{epochs}_epochs_conv.pt")
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,41 @@
+import os
+import sys
+import unittest
+import torch
+import torch.optim as optim
+import math
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing
+
+
+class TestPtoLrSchedulerCosineAnnealing(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.node = PtoLrSchedulerCosineAnnealing()
+
+        # Create a pseudo optimizer with learning rate parameter
+        self.model = torch.nn.Linear(2, 2)  # Dummy model
+        self.optimizer = optim.SGD(self.model.parameters(), lr=0.1)
+        self.num_epochs = 10
+        self.minimum_lr = 0.01
+
+    def test_lr_scheduler_cosine_annealing(self):
+        """Test if CosineAnnealingLR correctly updates learning rate."""
+        scheduler, = self.node.f(self.optimizer, self.num_epochs, self.minimum_lr)
+
+        # Ensure returned object is CosineAnnealingLR
+        self.assertIsInstance(scheduler, optim.lr_scheduler.CosineAnnealingLR, "Returned object is not a CosineAnnealingLR scheduler")
+
+        initial_lr = self.optimizer.param_groups[0]['lr']
+        for t in range(1, self.num_epochs + 1):
+            scheduler.step()
+            updated_lr = self.optimizer.param_groups[0]['lr']
+            expected_lr = self.minimum_lr + 0.5 * (initial_lr - self.minimum_lr) * (1 + math.cos(math.pi * t / self.num_epochs))
+            self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr} at step {t}")
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,65 @@
+import os
+import sys
+import unittest
+import torch
+import torch.optim as optim
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau  # Updated import
+
+
+class TestPtoLrSchedulerReduceOnPlateau(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.node = PtoLrSchedulerReduceOnPlateau()
+
+        # Create a pseudo optimizer with learning rate parameter
+        self.model = torch.nn.Linear(2, 2)  # Dummy model
+        self.optimizer = optim.SGD(self.model.parameters(), lr=0.1)
+        self.grace_period = 3
+        self.gamma = 0.5
+
+    def test_lr_scheduler_reduce_on_plateau(self):
+        """Test if ReduceLROnPlateau correctly updates learning rate."""
+        scheduler, = self.node.f(self.optimizer, self.grace_period, self.gamma)
+
+        # Ensure returned object is ReduceLROnPlateau
+        self.assertIsInstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau, "Returned object is not a ReduceLROnPlateau scheduler")
+
+        initial_lr = self.optimizer.param_groups[0]['lr']
+        # No improvement in loss
+        # Grace period map:   0,   1    2    3
+        # First epoch is the baseline, so it does not count.
+        # In this test case, the number 4 (the 4th) bad loss is greater than grace_period=3,
+        # so the loss reduction happens there.
+        loss_values = [0.5, 0.5, 0.5, 0.5, 0.5]  
+        for loss in loss_values:
+            scheduler.step(loss)
+        updated_lr = self.optimizer.param_groups[0]['lr']
+
+        # Check if learning rate was updated correctly after grace period
+        expected_lr = initial_lr * self.gamma
+        self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}")
+
+    def test_lr_scheduler_reduce_on_plateau_no_apply(self):
+        """Test if ReduceLROnPlateau correctly does not prematurely update learning rate."""
+        scheduler, = self.node.f(self.optimizer, self.grace_period, self.gamma)
+
+        # Ensure returned object is ReduceLROnPlateau
+        self.assertIsInstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau, "Returned object is not a ReduceLROnPlateau scheduler")
+
+        initial_lr = self.optimizer.param_groups[0]['lr']
+        loss_values = [0.5, 0.5, 0.5, 0.5]  
+        for loss in loss_values:
+            scheduler.step(loss)
+        updated_lr = self.optimizer.param_groups[0]['lr']
+
+        expected_lr = initial_lr
+        self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}")
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,44 @@
+import os
+import sys
+import unittest
+import torch
+import torch.optim as optim
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep  # Updated import
+
+
+class TestPtoLrSchedulerStep(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.node = PtoLrSchedulerStep()
+
+        # Create a pseudo optimizer with learning rate parameter
+        # For every 5 steps, if gamma is 0.1, then lr *= 0.5
+        # so 0.1 -> 0.05 -> 0.025
+        self.model = torch.nn.Linear(2, 2)  # Dummy model
+        self.optimizer = optim.SGD(self.model.parameters(), lr=0.1)
+        self.step_size = 5
+        self.gamma = 0.1
+
+    def test_lr_scheduler_step(self):
+        """Test if learning rate scheduler correctly updates learning rate."""
+        scheduler, = self.node.f(self.optimizer, self.step_size, self.gamma)
+
+        # Ensure returned object is StepLR
+        self.assertIsInstance(scheduler, optim.lr_scheduler.StepLR, "Returned object is not a StepLR scheduler")
+
+        initial_lr = self.optimizer.param_groups[0]['lr']
+        for _ in range(self.step_size):
+            scheduler.step()
+        updated_lr = self.optimizer.param_groups[0]['lr']
+
+        # Check if learning rate was updated correctly
+        expected_lr = initial_lr * self.gamma
+        self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}")
+
+if __name__ == "__main__":
+    unittest.main()