def export_model(self): """ Exports latest saved model to .nn format for Unity embedding. """ tf.train.write_graph(self.graph, self.cp_dir, 'raw_graph_def.pb', as_text=False) with self.graph.as_default(): target_nodes = ','.join(self._process_graph()) freeze_graph.freeze_graph( input_graph=os.path.join(self.cp_dir, 'raw_graph_def.pb'), input_binary=True, input_checkpoint=tf.train.latest_checkpoint(self.cp_dir), output_node_names=target_nodes, output_graph=(os.path.join(self.cp_dir, 'frozen_graph_def.pb')), # output_graph=(os.path.join(self.cp_dir,'model.bytes')), clear_devices=True, initializer_nodes='', input_saver='', restore_op_name='save/restore_all', filename_tensor_name='save/Const:0') tf2bc.convert(os.path.join(self.cp_dir, 'frozen_graph_def.pb'), os.path.join(self.cp_dir, 'model.nn'))
def export_policy_model(settings: SerializationSettings, graph: tf.Graph, sess: tf.Session) -> None: """ Exports latest saved model to .nn format for Unity embedding. """ frozen_graph_def = _make_frozen_graph(settings, graph, sess) # Save frozen graph frozen_graph_def_path = settings.model_path + "/frozen_graph_def.pb" with gfile.GFile(frozen_graph_def_path, "wb") as f: f.write(frozen_graph_def.SerializeToString()) # Convert to barracuda if settings.convert_to_barracuda: tf2bc.convert(frozen_graph_def_path, settings.model_path + ".nn") logger.info(f"Exported {settings.model_path}.nn file") # Save to onnx too (if we were able to import it) if ONNX_EXPORT_ENABLED and settings.convert_to_onnx: try: onnx_graph = convert_frozen_to_onnx(settings, frozen_graph_def) onnx_output_path = settings.model_path + ".onnx" with open(onnx_output_path, "wb") as f: f.write(onnx_graph.SerializeToString()) logger.info(f"Converting to {onnx_output_path}") except Exception: logger.exception( "Exception trying to save ONNX graph. Please report this error on " "https://github.com/Unity-Technologies/ml-agents/issues and " "attach a copy of frozen_graph_def.pb")
def export_policy_model( output_filepath: str, settings: SerializationSettings, graph: tf.Graph, sess: tf.Session, ) -> None: """ Exports a TF graph for a Policy to .nn and/or .onnx format for Unity embedding. :param output_filepath: file path to output the model (without file suffix) :param settings: SerializationSettings describing how to export the model :param graph: Tensorflow Graph for the policy :param sess: Tensorflow session for the policy """ frozen_graph_def = _make_frozen_graph(settings, graph, sess) if not os.path.exists(settings.model_path): os.makedirs(settings.model_path) # Save frozen graph frozen_graph_def_path = settings.model_path + "/frozen_graph_def.pb" with gfile.GFile(frozen_graph_def_path, "wb") as f: f.write(frozen_graph_def.SerializeToString()) # Convert to barracuda if settings.convert_to_barracuda: tf2bc.convert(frozen_graph_def_path, f"{output_filepath}.nn") logger.info(f"Exported {output_filepath}.nn") # Save to onnx too (if we were able to import it) if ONNX_EXPORT_ENABLED: if settings.convert_to_onnx: try: onnx_graph = convert_frozen_to_onnx(settings, frozen_graph_def) onnx_output_path = f"{output_filepath}.onnx" with open(onnx_output_path, "wb") as f: f.write(onnx_graph.SerializeToString()) logger.info(f"Converting to {onnx_output_path}") except Exception: # Make conversion errors fatal depending on environment variables (only done during CI) if _enforce_onnx_conversion(): raise logger.exception( "Exception trying to save ONNX graph. Please report this error on " "https://github.com/Unity-Technologies/ml-agents/issues and " "attach a copy of frozen_graph_def.pb") else: if _enforce_onnx_conversion(): raise RuntimeError( "ONNX conversion enforced, but couldn't import dependencies.")
def export_model(self): """ Exports latest saved model to .nn format for Unity embedding. """ with self.graph.as_default(): target_nodes = ",".join(self._process_graph()) graph_def = self.graph.as_graph_def() output_graph_def = graph_util.convert_variables_to_constants( self.sess, graph_def, target_nodes.replace(" ", "").split(",")) frozen_graph_def_path = self.model_path + "/frozen_graph_def.pb" with gfile.GFile(frozen_graph_def_path, "wb") as f: f.write(output_graph_def.SerializeToString()) tf2bc.convert(frozen_graph_def_path, self.model_path + ".nn") logger.info("Exported " + self.model_path + ".nn file")
def test_barracuda_converter(): path_prefix = os.path.dirname(os.path.abspath(__file__)) tmpfile = os.path.join(tempfile._get_default_tempdir(), next(tempfile._get_candidate_names()) + ".nn") # make sure there are no left-over files if os.path.isfile(tmpfile): os.remove(tmpfile) tf2bc.convert(path_prefix + "/BasicLearning.pb", tmpfile) # test if file exists after conversion assert os.path.isfile(tmpfile) # currently converter produces small output file even if input file is empty # 100 bytes is high enough to prove that conversion was successful assert os.path.getsize(tmpfile) > 100 # cleanup os.remove(tmpfile)
def export_policy_model( settings: SerializationSettings, graph: tf.Graph, sess: tf.Session ) -> None: """ Exports latest saved model to .nn format for Unity embedding. """ frozen_graph_def = _make_frozen_graph(settings, graph, sess) # Save frozen graph frozen_graph_def_path = settings.model_path + "/frozen_graph_def.pb" with gfile.GFile(frozen_graph_def_path, "wb") as f: f.write(frozen_graph_def.SerializeToString()) # Convert to barracuda if settings.convert_to_barracuda: tf2bc.convert(frozen_graph_def_path, settings.model_path + ".nn") logger.info(f"Exported {settings.model_path}.nn file") # Save to onnx too (if we were able to import it) if ONNX_EXPORT_ENABLED: if settings.convert_to_onnx: try: onnx_graph = convert_frozen_to_onnx(settings, frozen_graph_def) onnx_output_path = settings.model_path + ".onnx" with open(onnx_output_path, "wb") as f: f.write(onnx_graph.SerializeToString()) logger.info(f"Converting to {onnx_output_path}") except Exception: # Make conversion errors fatal depending on environment variables (only done during CI) if _enforce_onnx_conversion(): raise logger.exception( "Exception trying to save ONNX graph. Please report this error on " "https://github.com/Unity-Technologies/ml-agents/issues and " "attach a copy of frozen_graph_def.pb" ) else: if _enforce_onnx_conversion(): # Either we're on an old version of tensorflow, or the import failed. raise RuntimeError( "ONNX conversion enforced, but ONNX_EXPORT_ENABLED was false." )
def save(model, path): graph = model.graph print(graph) sess = model.sess # Make frozen graph with graph.as_default(): target_nodes = ",".join(_process_graph(graph)) graph_def = graph.as_graph_def() output_graph_def = graph_util.convert_variables_to_constants( sess, graph_def, target_nodes.replace(" ", "").split(",")) frozen_graph_def = output_graph_def # Useless step but I dont wanna change all the names # Save frozen graph frozen_graph_def_path = path + "/frozen_graph_def.pb" with gfile.GFile(frozen_graph_def_path, "wb") as f: f.write(frozen_graph_def.SerializeToString()) # Convert to barracuda tf2bc.convert(frozen_graph_def_path, path + ".nn")
def convert(input_cptk_file, output_nn_file): ########################################## # to flozen graph ########################################## #config = tf.ConfigProto() graph = tf.Graph() with tf.compat.v1.Session(graph=graph) as sess: # Restore from checkpoint loader = tf.compat.v1.train.import_meta_graph(input_cptk_file + '.meta') loader.restore(sess, input_cptk_file) # Create working directory shutil.rmtree(workdir) os.mkdir(workdir) # Export checkpoint to SavedModel builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(workdir) builder.add_meta_graph_and_variables( sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True) builder.save() # Convert frozen graph graph_def = graph.as_graph_def() output_graph_def = graph_util.convert_variables_to_constants( sess, graph_def, [ 'is_continuous_control', 'version_number', 'memory_size', 'action_output_shape', 'action_probs', 'action' ]) with gfile.GFile(workdir + "/frozen_graph_def.pb", "wb") as f: f.write(output_graph_def.SerializeToString()) ########################################## # to barracuda ########################################## tf2bc.convert(workdir + '/frozen_graph_def.pb', output_nn_file)
def export_model(self): """ Exports latest saved model to .nn format for Unity embedding. """ with self.graph.as_default(): target_nodes = ','.join(self._process_graph()) ckpt = tf.train.get_checkpoint_state(self.model_path) freeze_graph.freeze_graph( input_graph=self.model_path + '/raw_graph_def.pb', input_binary=True, input_checkpoint=ckpt.model_checkpoint_path, output_node_names=target_nodes, output_graph=(self.model_path + '/frozen_graph_def.pb'), clear_devices=True, initializer_nodes='', input_saver='', restore_op_name='save/restore_all', filename_tensor_name='save/Const:0') tf2bc.convert(self.model_path + '/frozen_graph_def.pb', self.model_path + '.nn') logger.info('Exported ' + self.model_path + '.nn file')