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')
