remove

ml-agents/mlagents/model_serialization.py

-from distutils.util import strtobool
-import os
-from typing import Any, List, Set, NamedTuple
-from distutils.version import LooseVersion
-
-try:
-    from tf2onnx.tfonnx import process_tf_graph, tf_optimize
-    from tf2onnx import optimizer
-
-    ONNX_EXPORT_ENABLED = True
-except ImportError:
-    # Either onnx and tf2onnx not installed, or they're not compatible with the version of tensorflow
-    ONNX_EXPORT_ENABLED = False
-    pass
-
-from mlagents.tf_utils import tf
-
-from tensorflow.python.platform import gfile
-from tensorflow.python.framework import graph_util
-
-from mlagents_envs.logging_util import get_logger
-from mlagents.trainers.tf import tensorflow_to_barracuda as tf2bc
-
-if LooseVersion(tf.__version__) < LooseVersion("1.12.0"):
-    # ONNX is only tested on 1.12.0 and later
-    ONNX_EXPORT_ENABLED = False
-
-logger = get_logger(__name__)
-
-
-POSSIBLE_INPUT_NODES = frozenset(
-    [
-        "action_masks",
-        "epsilon",
-        "prev_action",
-        "recurrent_in",
-        "sequence_length",
-        "vector_observation",
-    ]
-)
-
-POSSIBLE_OUTPUT_NODES = frozenset(
-    ["action", "action_probs", "recurrent_out", "value_estimate"]
-)
-
-MODEL_CONSTANTS = frozenset(
-    [
-        "action_output_shape",
-        "is_continuous_control",
-        "memory_size",
-        "version_number",
-        "trainer_major_version",
-        "trainer_minor_version",
-        "trainer_patch_version",
-    ]
-)
-VISUAL_OBSERVATION_PREFIX = "visual_observation_"
-
-
-class SerializationSettings(NamedTuple):
-    model_path: str
-    brain_name: str
-    convert_to_barracuda: bool = True
-    convert_to_onnx: bool = True
-    onnx_opset: int = 9
-
-
-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 _make_frozen_graph(
-    settings: SerializationSettings, graph: tf.Graph, sess: tf.Session
-) -> tf.GraphDef:
-    with graph.as_default():
-        target_nodes = ",".join(_process_graph(settings, graph))
-        graph_def = graph.as_graph_def()
-        output_graph_def = graph_util.convert_variables_to_constants(
-            sess, graph_def, target_nodes.replace(" ", "").split(",")
-        )
-    return output_graph_def
-
-
-def convert_frozen_to_onnx(
-    settings: SerializationSettings, frozen_graph_def: tf.GraphDef
-) -> Any:
-    # This is basically https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/convert.py
-
-    inputs = _get_input_node_names(frozen_graph_def)
-    outputs = _get_output_node_names(frozen_graph_def)
-    logger.info(f"onnx export - inputs:{inputs} outputs:{outputs}")
-
-    frozen_graph_def = tf_optimize(
-        inputs, outputs, frozen_graph_def, fold_constant=True
-    )
-
-    with tf.Graph().as_default() as tf_graph:
-        tf.import_graph_def(frozen_graph_def, name="")
-    with tf.Session(graph=tf_graph):
-        g = process_tf_graph(
-            tf_graph,
-            input_names=inputs,
-            output_names=outputs,
-            opset=settings.onnx_opset,
-        )
-
-    onnx_graph = optimizer.optimize_graph(g)
-    model_proto = onnx_graph.make_model(settings.brain_name)
-
-    return model_proto
-
-
-def _get_input_node_names(frozen_graph_def: Any) -> List[str]:
-    """
-    Get the list of input node names from the graph.
-    Names are suffixed with ":0"
-    """
-    node_names = _get_frozen_graph_node_names(frozen_graph_def)
-    input_names = node_names & POSSIBLE_INPUT_NODES
-
-    # Check visual inputs sequentially, and exit as soon as we don't find one
-    vis_index = 0
-    while True:
-        vis_node_name = f"{VISUAL_OBSERVATION_PREFIX}{vis_index}"
-        if vis_node_name in node_names:
-            input_names.add(vis_node_name)
-        else:
-            break
-        vis_index += 1
-    # Append the port
-    return [f"{n}:0" for n in input_names]
-
-
-def _get_output_node_names(frozen_graph_def: Any) -> List[str]:
-    """
-    Get the list of output node names from the graph.
-    Also include constants, so that they will be readable by the
-    onnx importer.
-    Names are suffixed with ":0"
-    """
-    node_names = _get_frozen_graph_node_names(frozen_graph_def)
-    output_names = node_names & (POSSIBLE_OUTPUT_NODES | MODEL_CONSTANTS)
-    # Append the port
-    return [f"{n}:0" for n in output_names]
-
-
-def _get_frozen_graph_node_names(frozen_graph_def: Any) -> Set[str]:
-    """
-    Get all the node names from the graph.
-    """
-    names = set()
-    for node in frozen_graph_def.node:
-        names.add(node.name)
-    return names
-
-
-def _process_graph(settings: SerializationSettings, graph: tf.Graph) -> List[str]:
-    """
-    Gets the list of the output nodes present in the graph for inference
-    :return: list of node names
-    """
-    all_nodes = [x.name for x in graph.as_graph_def().node]
-    nodes = [x for x in all_nodes if x in POSSIBLE_OUTPUT_NODES | MODEL_CONSTANTS]
-    logger.info("List of nodes to export for brain :" + settings.brain_name)
-    for n in nodes:
-        logger.info("\t" + n)
-    return nodes
-
-
-def _enforce_onnx_conversion() -> bool:
-    env_var_name = "TEST_ENFORCE_ONNX_CONVERSION"
-    if env_var_name not in os.environ:
-        return False
-
-    val = os.environ[env_var_name]
-    try:
-        # This handles e.g. "false" converting reasonably to False
-        return strtobool(val)
-    except Exception:
-        return False