浏览代码

onnx: export model constants as outputs, instead of in the initializer list (#4073)

* export model constants as outputs, instead of in the initializer list

* changelog
/MLA-1734-demo-provider
GitHub 5 年前
当前提交
43df1c8f
共有 2 个文件被更改,包括 5 次插入29 次删除
  1. 2
      com.unity.ml-agents/CHANGELOG.md
  2. 32
      ml-agents/mlagents/model_serialization.py

2
com.unity.ml-agents/CHANGELOG.md


### Bug Fixes
- Fixed an issue where SAC would perform too many model updates when resuming from a
checkpoint, and too few when using `buffer_init_steps`. (#4038)
- Fixed a bug in the onnx export that would cause constants needed for inference to not be visible to some versions of
the Barracuda importer. (#4073)
#### com.unity.ml-agents (C#)
#### ml-agents / ml-agents-envs / gym-unity (Python)

32
ml-agents/mlagents/model_serialization.py


from distutils.version import LooseVersion
try:
import onnx
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
from tf2onnx import optimizer

) -> Any:
# This is basically https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/convert.py
# Some constants in the graph need to be read by the inference system.
# These aren't used by the model anywhere, so trying to make sure they propagate
# through conversion and import is a losing battle. Instead, save them now,
# so that we can add them back later.
constant_values = {}
for n in frozen_graph_def.node:
if n.name in MODEL_CONSTANTS:
val = n.attr["value"].tensor.int_val[0]
constant_values[n.name] = val
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}")

onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model(settings.brain_name)
# Save the constant values back the graph initializer.
# This will ensure the importer gets them as global constants.
constant_nodes = []
for k, v in constant_values.items():
constant_node = _make_onnx_node_for_constant(k, v)
constant_nodes.append(constant_node)
model_proto.graph.initializer.extend(constant_nodes)
def _make_onnx_node_for_constant(name: str, value: int) -> Any:
tensor_value = onnx.TensorProto(
data_type=onnx.TensorProto.INT32,
name=name,
int32_data=[value],
dims=[1, 1, 1, 1],
)
return tensor_value
def _get_input_node_names(frozen_graph_def: Any) -> List[str]:

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.
output_names = node_names & POSSIBLE_OUTPUT_NODES
output_names = node_names & (POSSIBLE_OUTPUT_NODES | MODEL_CONSTANTS)
# Append the port
return [f"{n}:0" for n in output_names]

正在加载...
取消
保存