def test_load_graph(self): import google.protobuf as pb concrete_func = self.create_tf_function(concrete=True) with tmp_file(suffix=".pb") as path_pb, tmp_file( suffix=".pb.txt") as path_txt: cmsml.tensorflow.save_graph(path_txt, concrete_func, variables_to_constants=True) cmsml.tensorflow.save_graph(path_pb, concrete_func, variables_to_constants=False) self.assertTrue(os.path.exists(path_pb)) self.assertTrue(os.path.exists(path_txt)) graph = cmsml.tensorflow.load_graph(path_txt) self.assertIsInstance(graph, self.tf.Graph) graph = cmsml.tensorflow.load_graph(path_pb) self.assertIsInstance(graph, self.tf.Graph) with self.assertRaises(pb.text_format.ParseError): cmsml.tensorflow.load_graph(path_pb, as_text=True) with self.assertRaises(pb.message.DecodeError): cmsml.tensorflow.load_graph(path_txt, as_text=False)
def test_save_keras_model_v1(self): model = self.create_keras_model(self.tf1) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, model, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb.txt") as path: cmsml.tensorflow.save_graph(path, model, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, model, variables_to_constants=True) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, self.tf1.keras.backend.get_session(), variables_to_constants=False) self.assertTrue(os.path.exists(path))
def test_save_keras_model_v2(self): model = self.create_keras_model(self.tf) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, model, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, model, variables_to_constants=True) self.assertTrue(os.path.exists(path))
def test_save_empty_polymorphic_function(self): empty_poly_func = self.create_tf_function(no_input=True) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, empty_poly_func, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, empty_poly_func, variables_to_constants=True) self.assertTrue(os.path.exists(path))
def test_save_polymorphic_function_error(self): poly_func = self.create_tf_function() with self.assertRaises(ValueError): with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, poly_func, variables_to_constants=False) with self.assertRaises(ValueError): with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, poly_func, variables_to_constants=True)
def test_save_concrete_function(self): concrete_func = self.create_tf_function(concrete=True) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, concrete_func, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb.txt") as path: cmsml.tensorflow.save_graph(path, concrete_func, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, concrete_func, variables_to_constants=True) self.assertTrue(os.path.exists(path))
def test_save_graph(self): graph, session = self.create_tf1_graph() if graph is None or session is None: return with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, graph, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb.txt") as path: cmsml.tensorflow.save_graph(path, graph, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, graph.as_graph_def(), variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, session, variables_to_constants=False) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: cmsml.tensorflow.save_graph(path, session, variables_to_constants=True, output_names=["output"]) self.assertTrue(os.path.exists(path)) with tmp_file(suffix=".pb") as path: with self.assertRaises(ValueError): cmsml.tensorflow.save_graph(path, session, variables_to_constants=True) self.assertFalse(os.path.exists(path))
def test_write_summary(self): concrete_func = self.create_tf_function(concrete=True) with tmp_dir(create=False) as path: cmsml.tensorflow.write_graph_summary(concrete_func.graph, path) self.assertTrue(os.path.exists(path)) self.assertGreater(len(os.listdir(path)), 0) with tmp_file(suffix=".pb") as graph_path: cmsml.tensorflow.save_graph(graph_path, concrete_func) with tmp_dir(create=False) as path: cmsml.tensorflow.write_graph_summary(graph_path, path) self.assertTrue(os.path.exists(path)) self.assertGreater(len(os.listdir(path)), 0) self.assertTrue(os.path.exists(path))
def test_load_graph_and_run(self): import numpy as np tf = self.tf1 if tf is None: return _, session = self.create_tf1_graph() with tmp_file(suffix=".pb.txt") as path: cmsml.tensorflow.save_graph(path, session, variables_to_constants=True, output_names=["output"]) graph = cmsml.tensorflow.load_graph(path) session = self.create_tf1_session(graph) with graph.as_default(): x = graph.get_tensor_by_name("input:0") y = graph.get_tensor_by_name("output:0") out = session.run(y, {x: np.ones((2, 10))}) self.assertEqual(out.shape, (2, 1)) self.assertEqual(tuple(out[..., 0]), (1., 1.))