Add some tests

metrics/test_aggregation.py

@@ -0,0 +1,42 @@
+import torch
+import unittest
+from aggregation import StableMean
+
+
+class TestStableMean(unittest.TestCase):
+
+    def setUp(self):
+        self.metric = StableMean()
+
+    def test_compute_empty(self):
+        result = self.metric.compute()
+        self.assertEqual(result, torch.tensor(0.0))
+
+    def test_compute_single_value(self):
+        self.metric.update(torch.tensor(1.0))
+        result = self.metric.compute()
+        self.assertEqual(result, torch.tensor(1.0))
+
+    def test_compute_weighted_single_value(self):
+        self.metric.update(torch.tensor(1.0), weight=torch.tensor(2.0))
+        result = self.metric.compute()
+        self.assertEqual(result, torch.tensor(1.0))
+
+    def test_compute_multiple_values(self):
+        self.metric.update(torch.tensor(1.0))
+        self.metric.update(torch.tensor(2.0))
+        self.metric.update(torch.tensor(3.0))
+        result = self.metric.compute()
+        self.assertEqual(result, torch.tensor(2.0))
+
+    def test_compute_weighted_multiple_values(self):
+        self.metric.update(torch.tensor(1.0), weight=torch.tensor(1.0))
+        self.metric.update(torch.tensor(2.0), weight=torch.tensor(2.0))
+        self.metric.update(torch.tensor(3.0), weight=torch.tensor(3.0))
+        result = self.metric.compute()
+        print(f"get=  {result.item()} but expected= 2.1666666667")
+        self.assertAlmostEqual(result.item(), 2.1666666667, places=0)
+
+
+if '__name__' == '__main__':
+    unittest.main()

optimizers/test_optimizer.py

@@ -0,0 +1,44 @@
+import unittest
+import torch
+from tml.optimizers.config import LearningRate, OptimizerConfig
+from optimizer import compute_lr, LRShim, get_optimizer_class, build_optimizer
+
+
+class TestComputeLR(unittest.TestCase):
+    def test_constant_lr(self):
+        lr_config = LearningRate(constant=0.1)
+        lr = compute_lr(lr_config, step=0)
+        self.assertAlmostEqual(lr, 0.1)
+
+    def test_piecewise_constant_lr(self):
+        lr_config = LearningRate(piecewise_constant={"learning_rate_boundaries": [10, 20], "learning_rate_values": [0.1, 0.01, 0.001]})
+        lr = compute_lr(lr_config, step=5)
+        self.assertAlmostEqual(lr, 0.1)
+        lr = compute_lr(lr_config, step=15)
+        self.assertAlmostEqual(lr, 0.01)
+        lr = compute_lr(lr_config, step=25)
+        self.assertAlmostEqual(lr, 0.001)
+
+
+class TestLRShim(unittest.TestCase):
+    def setUp(self):
+        self.optimizer = torch.optim.SGD([torch.randn(10, 10)], lr=0.1)
+        self.lr_dict = {"ALL_PARAMS": LearningRate(constant=0.1)}
+
+    def test_get_lr(self):
+        lr_scheduler = LRShim(self.optimizer, self.lr_dict)
+        lr = lr_scheduler.get_lr()
+        self.assertAlmostEqual(lr, [0.1])
+
+
+class TestBuildOptimizer(unittest.TestCase):
+    def test_build_optimizer(self):
+        model = torch.nn.Linear(10, 1)
+        optimizer_config = OptimizerConfig(sgd={"lr": 0.1})
+        optimizer, scheduler = build_optimizer(model, optimizer_config)
+        self.assertIsInstance(optimizer, torch.optim.SGD)
+        self.assertIsInstance(scheduler, LRShim)
+
+
+if __name__ == "__main__":
+    unittest.main()

reader/utils.py

@@ -21,8 +21,7 @@ def roundrobin(*iterables):
   while num_active:
     try:
       for _next in nexts:
-        result = _next()
-        yield result
+        yield _next()
     except StopIteration:
       # Remove the iterator we just exhausted from the cycle.
       num_active -= 1