class TestFunctionSet(TestCase): def setUp(self): self.fs = FunctionSet( a = Linear(3, 2), b = Linear(3, 2) ) def test_get_sorted_funcs(self): self.assertItemsEqual([k for (k, v) in self.fs._get_sorted_funcs()], ('a', 'b')) def check_equal_fs(self, fs1, fs2): self.assertTrue((fs1.a.W == fs2.a.W).all()) self.assertTrue((fs1.a.b == fs2.a.b).all()) self.assertTrue((fs1.b.W == fs2.b.W).all()) self.assertTrue((fs1.b.b == fs2.b.b).all()) def test_pickle_cpu(self): s = pickle.dumps(self.fs) fs2 = pickle.loads(s) self.check_equal_fs(self.fs, fs2) @attr.gpu def test_pickle_gpu(self): self.fs.to_gpu() s = pickle.dumps(self.fs) fs2 = pickle.loads(s) self.fs.to_cpu() fs2.to_cpu() self.check_equal_fs(self.fs, fs2)
class TestFunctionSet(TestCase): def setUp(self): self.fs = FunctionSet(a=Linear(3, 2), b=Linear(3, 2)) def test_get_sorted_funcs(self): assertCountEqual(self, [k for (k, v) in self.fs._get_sorted_funcs()], ('a', 'b')) def check_equal_fs(self, fs1, fs2): self.assertTrue((fs1.a.W == fs2.a.W).all()) self.assertTrue((fs1.a.b == fs2.a.b).all()) self.assertTrue((fs1.b.W == fs2.b.W).all()) self.assertTrue((fs1.b.b == fs2.b.b).all()) def test_pickle_cpu(self): s = pickle.dumps(self.fs) fs2 = pickle.loads(s) self.check_equal_fs(self.fs, fs2) @attr.gpu def test_pickle_gpu(self): self.fs.to_gpu() s = pickle.dumps(self.fs) fs2 = pickle.loads(s) self.fs.to_cpu() fs2.to_cpu() self.check_equal_fs(self.fs, fs2)
class TestNestedFunctionSet(TestCase): def setUp(self): self.fs1 = FunctionSet( a = MockFunction((1, 2))) self.fs2 = FunctionSet( fs1 = self.fs1, b = MockFunction((3, 4))) def test_get_sorted_funcs(self): self.assertItemsEqual([k for (k, v) in self.fs2._get_sorted_funcs()], ('b', 'fs1')) def test_collect_parameters(self): p_b = np.zeros((3, 4)).astype(np.float32) p_a = np.zeros((1, 2)).astype(np.float32) gp_b = np.ones((3, 4)).astype(np.float32) gp_a = np.ones((1, 2)).astype(np.float32) actual = self.fs2.collect_parameters() self.assertTrue(map(len, actual) == [2, 2]) self.assertTrue((actual[0][0] == p_b).all()) self.assertTrue((actual[0][1] == p_a).all()) self.assertTrue((actual[1][0] == gp_b).all()) self.assertTrue((actual[1][1] == gp_a).all()) def test_pickle_cpu(self): fs2_serialized = pickle.dumps(self.fs2) fs2_loaded = pickle.loads(fs2_serialized) self.assertTrue((self.fs2.b.p == fs2_loaded.b.p).all()) self.assertTrue((self.fs2.fs1.a.p == fs2_loaded.fs1.a.p).all()) @attr.gpu def test_pickle_gpu(self): self.fs2.to_gpu() fs2_serialized = pickle.dumps(self.fs2) fs2_loaded = pickle.loads(fs2_serialized) fs2_loaded.to_cpu() self.fs2.to_cpu() self.assertTrue((self.fs2.b.p == fs2_loaded.b.p).all()) self.assertTrue((self.fs2.fs1.a.p == fs2_loaded.fs1.a.p).all())
class TestNestedFunctionSet(TestCase): def setUp(self): self.fs1 = FunctionSet(a=MockFunction((1, 2))) self.fs2 = FunctionSet(fs1=self.fs1, b=MockFunction((3, 4))) def test_get_sorted_funcs(self): assertCountEqual(self, [k for (k, v) in self.fs2._get_sorted_funcs()], ('b', 'fs1')) def test_collect_parameters(self): p_b = np.zeros((3, 4)).astype(np.float32) p_a = np.zeros((1, 2)).astype(np.float32) gp_b = np.ones((3, 4)).astype(np.float32) gp_a = np.ones((1, 2)).astype(np.float32) actual = self.fs2.collect_parameters() self.assertTrue(list(map(len, actual)) == [2, 2]) self.assertTrue((actual[0][0] == p_b).all()) self.assertTrue((actual[0][1] == p_a).all()) self.assertTrue((actual[1][0] == gp_b).all()) self.assertTrue((actual[1][1] == gp_a).all()) def test_pickle_cpu(self): fs2_serialized = pickle.dumps(self.fs2) fs2_loaded = pickle.loads(fs2_serialized) self.assertTrue((self.fs2.b.p == fs2_loaded.b.p).all()) self.assertTrue((self.fs2.fs1.a.p == fs2_loaded.fs1.a.p).all()) @attr.gpu def test_pickle_gpu(self): self.fs2.to_gpu() fs2_serialized = pickle.dumps(self.fs2) fs2_loaded = pickle.loads(fs2_serialized) fs2_loaded.to_cpu() self.fs2.to_cpu() self.assertTrue((self.fs2.b.p == fs2_loaded.b.p).all()) self.assertTrue((self.fs2.fs1.a.p == fs2_loaded.fs1.a.p).all())