Rewrite tests in unittest framework

Author: BillHuang2001

unit_test/algorithms/test_pso.py

@@ -1,58 +1,55 @@
+from unittest import TestCase
 import time
 import torch
 from torch.profiler import profile, ProfilerActivity
+from evox.core import vmap, Problem, use_state, jit
+from evox.workflows import StdWorkflow
+from evox.algorithms import PSO
 
-import os
-import sys
-current_directory = os.getcwd()
-if current_directory not in sys.path:
-    sys.path.append(current_directory)
 
-from src.core import vmap, Problem, use_state, jit
-from src.workflows import StdWorkflow 
-from src.algorithms import PSO
+class TestPSO(TestCase):
+    def setUp(self):
+        class Sphere(Problem):
+            def __init__(self):
+                super().__init__()
 
+            def evaluate(self, pop: torch.Tensor):
+                return (pop**2).sum(-1)
 
-if __name__ == "__main__":
-    class Sphere(Problem):
-        def __init__(self):
-            super().__init__()
+        self.algo = PSO(pop_size=10, lb=-10 * torch.ones(3), ub=10 * torch.ones(3))
+        self.prob = Sphere()
 
-        def evaluate(self, pop: torch.Tensor):
-            return (pop**2).sum(-1)
-
-    torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
-    print(torch.get_default_device())
-    algo = PSO(pop_size=10, lb=-10 * torch.ones(3), ub=10 * torch.ones(3))
-    prob = Sphere()
-    workflow = StdWorkflow()
-    workflow.setup(algo, prob)
-    workflow.init_step()
-    workflow.step()
-    # with open("tests/a.md", "w") as ff:
-    #     ff.write(workflow.step.inlined_graph.__str__())
-    state_step = use_state(lambda: workflow.step)
-    vmap_state_step = vmap(state_step)
-    print(vmap_state_step.init_state(2))
-    state = state_step.init_state()
-    jit_state_step = jit(state_step, trace=True, example_inputs=(state,))
-    state = state_step.init_state()
-    # with open("tests/b.md", "w") as ff:
-    #     ff.write(jit_state_step.inlined_graph.__str__())
-    t = time.time()
-    with profile(
-        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, profile_memory=True
-    ) as prof:
-        for _ in range(1000):
-            workflow.step()
-    print(prof.key_averages().table())
-    torch.cuda.synchronize()
-    t = time.time()
-    with profile(
-        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, profile_memory=True
-    ) as prof:
-        for _ in range(1000):
-            state = jit_state_step(state)
-    print(prof.key_averages().table())
-    torch.cuda.synchronize()
-    print(time.time() - t)
+    def test_pso(self):
+        torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
+        print(torch.get_default_device())
+        workflow = StdWorkflow()
+        workflow.setup(self.algo, self.prob)
+        workflow.init_step()
+        workflow.step()
+        state_step = use_state(lambda: workflow.step)
+        vmap_state_step = vmap(state_step)
+        print(vmap_state_step.init_state(2))
+        state = state_step.init_state()
+        jit_state_step = jit(state_step, trace=True, example_inputs=(state,))
+        state = state_step.init_state()
+        t = time.time()
+        with profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            record_shapes=True,
+            profile_memory=True,
+        ) as prof:
+            for _ in range(1000):
+                workflow.step()
+        print(prof.key_averages().table())
+        torch.cuda.synchronize()
+        t = time.time()
+        with profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            record_shapes=True,
+            profile_memory=True,
+        ) as prof:
+            for _ in range(1000):
+                state = jit_state_step(state)
+        print(prof.key_averages().table())
+        torch.cuda.synchronize()
+        print(time.time() - t)

unit_test/core/test_jit_util.py

@@ -1,21 +1,28 @@
+import unittest
 from functools import partial
 import torch
+from evox.core import vmap, jit
 
-import os
-import sys
-current_directory = os.getcwd()
-if current_directory not in sys.path:
-    sys.path.append(current_directory)
 
-from src.core import vmap, jit
+@partial(vmap, example_ndim=2)
+def _single_eval(
+    x: torch.Tensor, p: float = 2.0, q: torch.Tensor = torch.as_tensor(range(2))
+):
+    return (x**p).sum() * q.sum()
 
 
-if __name__ == "__main__":
+class TestJitUtil(unittest.TestCase):
+    def setUp(self):
+        self.expected = torch.tensor([8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0])
 
-    @partial(vmap, example_ndim=2)
-    def _single_eval(x: torch.Tensor, p: float = 2.0, q: torch.Tensor = torch.as_tensor(range(2))):
-        return (x**p).sum() * q.sum()
+    def test_single_eval(self):
+        result = _single_eval(2 * torch.ones(10, 2))
+        self.assertTrue(torch.equal(result, self.expected))
 
-    print(_single_eval(2 * torch.ones(10, 2)))
-    print(jit(_single_eval)(2 * torch.ones(10, 2)))
-    print(jit(_single_eval, trace=True, lazy=True)(2 * torch.ones(10, 2)))
+    def test_jit_single_eval(self):
+        result = jit(_single_eval)(2 * torch.ones(10, 2))
+        self.assertTrue(torch.equal(result, self.expected))
+
+    def test_jit_single_eval_trace_lazy(self):
+        result = jit(_single_eval, trace=True, lazy=True)(2 * torch.ones(10, 2))
+        self.assertTrue(torch.equal(result, self.expected))

unit_test/core/test_module.py

@@ -1,69 +1,97 @@
+import unittest
 import torch
 import torch.nn as nn
 from typing import Dict, List
 
-import os
-import sys
-current_directory = os.getcwd()
-if current_directory not in sys.path:
-    sys.path.append(current_directory)
-
-from src.core import jit_class, ModuleBase, trace_impl, use_state
-
-
-if __name__ == "__main__":
-
-    @jit_class
-    class Test(ModuleBase):
-
-        def __init__(self, threshold=0.5):
-            super().__init__()
-            self.threshold = threshold
-            self.sub_mod = nn.Module()
-            self.sub_mod.buf = nn.Buffer(torch.zeros(()))
-
-        def h(self, q: torch.Tensor) -> torch.Tensor:
-            if q.flatten()[0] > self.threshold:
-                x = torch.sin(q)
-            else:
-                x = torch.tan(q)
-            return x * x.shape[1]
-
-        @trace_impl(h)
-        def th(self, q: torch.Tensor) -> torch.Tensor:
-            x = torch.where(q.flatten()[0] > self.threshold, q + 2, q + 5)
-            x += self.g(x).abs()
-            x *= x.shape[1]
-            self.sub_mod.buf = x.sum()
-            return x
-
-        def g(self, p: torch.Tensor) -> torch.Tensor:
-            x = torch.cos(p)
-            return x * p.shape[0]
-
-    t = Test()
-    print(t.h.inlined_graph)
-    result = t.g(torch.rand(100, 4))
-    print(result)
-    t.add_mutable("mut_list", [torch.zeros(10), torch.ones(10)])
-    t.add_mutable("mut_dict", {"a": torch.zeros(20), "b": torch.ones(20)})
-    print(t.mut_list[0])
-    print(t.mut_dict["b"])
-
-    t = Test()
-    fn = use_state(lambda: t.h, is_generator=True)
-    trace_fn = torch.jit.trace(fn, (fn.init_state(), torch.ones(10, 1)), strict=False)
-
-    def loop(init_state: Dict[str, torch.Tensor], init_x: torch.Tensor, n: int = 10):
-        state = init_state
-        ret = init_x
-        rets: List[torch.Tensor] = []
-        for _ in range(n):
-            state, ret = trace_fn(state, ret)
-            rets.append(state["self.sub_mod.buf"])
-        return rets
-
-    print(trace_fn.code)
-    loop = torch.jit.trace(loop, (fn.init_state(), torch.rand(10, 2)), strict=False)
-    print(loop.code)
-    print(loop(fn.init_state(), torch.rand(10, 2)))
+from evox.core import jit_class, ModuleBase, trace_impl, use_state
+
+
+@jit_class
+class DummyModule(ModuleBase):
+
+    def __init__(self, threshold=0.5):
+        super().__init__()
+        self.threshold = threshold
+        self.sub_mod = nn.Module()
+        self.sub_mod.buf = nn.Buffer(torch.zeros(()))
+
+    def h(self, q: torch.Tensor) -> torch.Tensor:
+        if q.flatten()[0] > self.threshold:
+            x = torch.sin(q)
+        else:
+            x = torch.tan(q)
+        return x * x.shape[1]
+
+    @trace_impl(h)
+    def th(self, q: torch.Tensor) -> torch.Tensor:
+        x = torch.where(q.flatten()[0] > self.threshold, q + 2, q + 5)
+        x += self.g(x).abs()
+        x *= x.shape[1]
+        self.sub_mod.buf = x.sum()
+        return x
+
+    def g(self, p: torch.Tensor) -> torch.Tensor:
+        x = torch.cos(p)
+        return x * p.shape[0]
+
+
+class TestModule(unittest.TestCase):
+
+    def setUp(self):
+        self.test_instance = DummyModule()
+
+    def test_h_function(self):
+        q = torch.rand(100, 4)
+        result = self.test_instance.h(q)
+        self.assertIsInstance(result, torch.Tensor)
+
+    def test_g_function(self):
+        p = torch.rand(100, 4)
+        result = self.test_instance.g(p)
+        self.assertIsInstance(result, torch.Tensor)
+
+    def test_th_function(self):
+        q = torch.rand(100, 4)
+        result = self.test_instance.th(q)
+        self.assertIsInstance(result, torch.Tensor)
+
+    def test_add_mutable_list(self):
+        self.test_instance.add_mutable("mut_list", [torch.zeros(10), torch.ones(10)])
+        self.assertTrue(torch.equal(self.test_instance.mut_list[0], torch.zeros(10)))
+        self.assertTrue(torch.equal(self.test_instance.mut_list[1], torch.ones(10)))
+
+    def test_add_mutable_dict(self):
+        self.test_instance.add_mutable(
+            "mut_dict", {"a": torch.zeros(20), "b": torch.ones(20)}
+        )
+        self.assertTrue(torch.equal(self.test_instance.mut_dict["a"], torch.zeros(20)))
+        self.assertTrue(torch.equal(self.test_instance.mut_dict["b"], torch.ones(20)))
+
+    def test_trace_fn(self):
+        fn = use_state(lambda: self.test_instance.h, is_generator=True)
+        trace_fn = torch.jit.trace(
+            fn, (fn.init_state(), torch.ones(10, 1)), strict=False
+        )
+        self.assertIsNotNone(trace_fn)
+
+    def test_loop_function(self):
+        fn = use_state(lambda: self.test_instance.h, is_generator=True)
+        trace_fn = torch.jit.trace(
+            fn, (fn.init_state(), torch.ones(10, 1)), strict=False
+        )
+
+        def loop(
+            init_state: Dict[str, torch.Tensor], init_x: torch.Tensor, n: int = 10
+        ):
+            state = init_state
+            ret = init_x
+            rets: List[torch.Tensor] = []
+            for _ in range(n):
+                state, ret = trace_fn(state, ret)
+                rets.append(state["self.sub_mod.buf"])
+            return rets
+
+        loop_traced = torch.jit.trace(
+            loop, (fn.init_state(), torch.rand(10, 2)), strict=False
+        )
+        self.assertIsNotNone(loop_traced)