Skip to content

SeetaResearch / Dragon

Go to a project
Commit b93bde0d by Ting PAN
Options

Refactor ONNX frontends and backends 

Summary:
This commit redesigns the ``vm.onnx`` by referring the official repository.
Frontends and backends are aligned with identical API for dragon, torch and tensorrt.
1 parent e82d2ba4
Inline Side-by-side
Showing with 1023 additions and 861 deletions
docs/api/python/onnx.rst → docs/api/python/dragon/onnx.rst
vm.onnx
=======
dragon.onnx
===========
.. toctree::
:hidden:
.. only:: html
Classes
-------
`class BackendRep <onnx/BackendRep.html>`_
: ONNX-Dragon backend to execute repeatedly.
Functions
---------
`export(...) <onnx/export.html>`_
: Export the recorded graph to an onnx model.
onnx/Shell
`prepare_backend(...) <onnx/prepare_backend.html>`_
: Create a backend to execute repeatedly.
`record(...) <onnx/record.html>`_
: Context-manger to record the graph.
`run_model(...) <onnx/run_model.html>`_
: Execute an onnx model once.
`supports_device(...) <onnx/supports_device.html>`_
: Query if the given device is supported to execute.
Operators
#########
---------
======================== ========= ========================================
Name Supported Reference
......@@ -168,13 +189,15 @@ Name Supported Reference
`Xor`_ |v| :func:`dragon.bitwise.bitwise_xor`
======================== ========= ========================================
.. only:: html
Classes
#######
.. toctree::
:hidden:
`class Shell <onnx/Shell.html>`_
: Context-manger to export or load onnx models.
onnx/BackendRep
onnx/prepare_backend
onnx/export
onnx/record
onnx/run_model
onnx/supports_device
.. _Abs: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Abs
.. _Acos: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Acos
......@@ -331,14 +354,14 @@ Name Supported Reference
.. _Where: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Where
.. _Xor: https://github.com/onnx/onnx/blob/master/docs/Operators.md#Xor
.. |v| image:: ../_static/images/tick.png
.. |v| image:: ../../_static/images/tick.png
:height: 18
.. raw:: html
<style>
h1:before {
content: "Module: dragon.";
content: "Module: ";
color: #103d3e;
}
</style>
docs/api/python/onnx/Shell.rst → docs/api/python/dragon/onnx/BackendRep.rst
Shell
=====
BackendRep
==========
.. autoclass:: dragon.vm.onnx.Shell
.. autoclass:: dragon.onnx.BackendRep
__init__
--------
.. automethod:: dragon.vm.onnx.Shell.__init__
.. automethod:: dragon.onnx.BackendRep.__init__
Methods
-------
as_default
##########
.. automethod:: dragon.vm.onnx.Shell.as_default
export
######
.. automethod:: dragon.vm.onnx.Shell.export
load_model
##########
.. automethod:: dragon.vm.onnx.Shell.load_model
run
###
.. automethod:: dragon.onnx.BackendRep.run
.. raw:: html
......
docs/api/python/tensorrt/backend/prepare.rst → docs/api/python/dragon/onnx/export.rst
prepare
=======
export
======
.. autofunction:: dragon.vm.tensorrt.backend.prepare
.. autofunction:: dragon.onnx.export
.. raw:: html
<style>
h1:before {
content: "tensorrt.backend.";
content: "dragon.onnx.";
color: #103d3e;
}
</style>
docs/api/python/dragon/onnx/prepare_backend.rst 0 → 100644
prepare_backend
===============
.. autofunction:: dragon.onnx.prepare_backend
.. raw:: html
<style>
h1:before {
content: "dragon.onnx.";
color: #103d3e;
}
</style>
docs/api/python/dragon/onnx/record.rst 0 → 100644
record
======
.. autofunction:: dragon.onnx.record
.. _dragon.onnx.export(...): export.html
.. raw:: html
<style>
h1:before {
content: "dragon.onnx.";
color: #103d3e;
}
</style>
docs/api/python/tensorrt/backend/run_model.rst → docs/api/python/dragon/onnx/run_model.rst
run_model
=========
.. autofunction:: dragon.vm.tensorrt.backend.run_model
.. autofunction:: dragon.onnx.run_model
.. raw:: html
<style>
h1:before {
content: "tensorrt.backend.";
content: "dragon.onnx.";
color: #103d3e;
}
</style>
docs/api/python/tensorrt/backend/supports_device.rst → docs/api/python/dragon/onnx/supports_device.rst
supports_device
===============
.. autofunction:: dragon.vm.tensorrt.backend.supports_device
.. autofunction:: dragon.onnx.supports_device
.. raw:: html
<style>
h1:before {
content: "tensorrt.backend.";
content: "dragon.onnx.";
color: #103d3e;
}
</style>
docs/api/python/index.rst
......@@ -38,6 +38,7 @@ Dragon
* `dragon.math <dragon/math.html>`_
* `dragon.metrics <dragon/metrics.html>`_
* `dragon.nn <dragon/nn.html>`_
* `dragon.onnx <dragon/onnx.html>`_
* `dragon.optimizers <dragon/optimizers.html>`_
* `dragon.random <dragon/random.html>`_
* `dragon.vision <dragon/vision.html>`_
......@@ -146,7 +147,9 @@ ONNX
This integration involves the following components:
* `onnx <onnx.html>`_
* `dragon.onnx <dragon/onnx.html>`_
* `tensorrt.onnx <tensorrt/onnx.html>`_
* `torch.onnx <torch/onnx.html>`_
TensorRT
########
......@@ -158,7 +161,7 @@ TensorRT
This integration involves the following components:
* `tensorrt <tensorrt.html>`_
* `tensorrt.backend <tensorrt/backend.html>`_
* `tensorrt.onnx <tensorrt/onnx.html>`_
Modules
-------
......@@ -198,6 +201,9 @@ Modules
`Module nn <dragon/nn.html>`_
: Native API for ``dragon.nn`` namespace.
`Module onnx <dragon/onnx.html>`_
: Native API for ``dragon.onnx`` namespace.
`Module optimizers <dragon/optimizers.html>`_
: Native API for ``dragon.optimizers`` namespace.
......@@ -219,9 +225,6 @@ Modules
`Module vm.dali.ops <dali/ops.html>`_
: Virtual API for ``dali.ops`` namespace.
`Module vm.onnx <onnx.html>`_
: Virtual API for ``onnx`` namespace.
`Module vm.tensorflow <tensorflow.html>`_
: Virtual API for ``tensorflow`` namespace.
......@@ -255,8 +258,11 @@ Modules
`Module vm.tensorlayer.models <tensorlayer/models.html>`_
: Virtual API for ``tensorlayer.models`` namespace.
`Module vm.tensorrt.backend <tensorrt/backend.html>`_
: Virtual API for ``tensorrt.backend`` namespace.
`Module vm.tensorrt <tensorrt.html>`_
: Virtual API for ``tensorrt`` namespace.
`Module vm.tensorrt.onnx <tensorrt/onnx.html>`_
: Virtual API for ``tensorrt.onnx`` namespace.
`Module vm.torch <torch.html>`_
: Virtual API for ``torch`` namespace.
......@@ -306,6 +312,7 @@ Modules
dragon/math
dragon/metrics
dragon/nn
dragon/onnx
dragon/optimizers
dragon/random
dragon/vision
......@@ -313,7 +320,6 @@ Modules
caffe/layers
dali
dali/ops
onnx
tensorflow
tensorflow/bitwise
tensorflow/dtypes
......@@ -326,7 +332,7 @@ Modules
tensorlayer/layers
tensorlayer/models
tensorrt
tensorrt/backend
tensorrt/onnx
torch
torch/autograd
torch/distributed
......
docs/api/python/tensorrt.rst
......@@ -12,15 +12,11 @@ vm.tensorrt
`class Engine <tensorrt/Engine.html>`_
: The executing engine with bindings.
`class ONNXBackendRep <tensorrt/ONNXBackendRep.html>`_
: Load and run onnx models.
.. toctree::
:hidden:
tensorrt/Binding
tensorrt/Engine
tensorrt/ONNXBackendRep
.. raw:: html
......
docs/api/python/tensorrt/backend.rst → docs/api/python/tensorrt/onnx.rst
vm.tensorrt.backend
===================
vm.tensorrt.onnx
================
.. only:: html
Classes
-------
`class BackendRep <onnx/BackendRep.html>`_
: ONNX-TensorRT backend to execute repeatedly.
Functions
---------
`prepare(...) <backend/prepare.html>`_
: Build a TensorRT engine from the onnx model.
`prepare_backend(...) <onnx/prepare_backend.html>`_
: Create a backend to execute repeatedly.
`run_model(...) <backend/run_model.html>`_
: Build and run a TensorRT engine from the onnx model.
: Execute an onnx model once.
`run_node(...) <backend/run_node.html>`_
: Build and run a TensorRT engine from the onnx node.
: Execute an onnx node once.
`supports_device(...) <backend/supports_device.html>`_
: Query if given device is supported.
: Query if the given device is supported to execute.
.. toctree::
:hidden:
backend/prepare
backend/run_model
backend/run_node
backend/supports_device
onnx/BackendRep
onnx/prepare
onnx/run_model
onnx/run_node
onnx/supports_device
.. raw:: html
......
docs/api/python/tensorrt/ONNXBackendRep.rst → docs/api/python/tensorrt/onnx/BackendRep.rst
ONNXBackendRep
==============
BackendRep
==========
.. autoclass:: dragon.vm.tensorrt.ONNXBackendRep
.. autoclass:: dragon.vm.tensorrt.onnx.BackendRep
__init__
--------
.. automethod:: dragon.vm.tensorrt.ONNXBackendRep.__init__
.. automethod:: dragon.vm.tensorrt.onnx.BackendRep.__init__
Properties
----------
engine
######
.. autoattribute:: dragon.vm.tensorrt.ONNXBackendRep.engine
.. autoattribute:: dragon.vm.tensorrt.onnx.BackendRep.engine
Methods
......@@ -20,7 +20,7 @@ Methods
run
###
.. automethod:: dragon.vm.tensorrt.ONNXBackendRep.run
.. automethod:: dragon.vm.tensorrt.onnx.BackendRep.run
.. raw:: html
......
docs/api/python/tensorrt/onnx/prepare.rst 0 → 100644
prepare_backend
===============
.. autofunction:: dragon.vm.tensorrt.onnx.prepare_backend
.. raw:: html
<style>
h1:before {
content: "tensorrt.onnx.";
color: #103d3e;
}
</style>
docs/api/python/tensorrt/onnx/run_model.rst 0 → 100644
run_model
=========
.. autofunction:: dragon.vm.tensorrt.onnx.run_model
.. raw:: html
<style>
h1:before {
content: "tensorrt.onnx.";
color: #103d3e;
}
</style>
docs/api/python/tensorrt/backend/run_node.rst → docs/api/python/tensorrt/onnx/run_node.rst
run_node
========
.. autofunction:: dragon.vm.tensorrt.backend.run_node
.. autofunction:: dragon.vm.tensorrt.onnx.run_node
.. raw:: html
<style>
h1:before {
content: "tensorrt.backend.";
content: "tensorrt.onnx.";
color: #103d3e;
}
</style>
docs/api/python/tensorrt/onnx/supports_device.rst 0 → 100644
supports_device
===============
.. autofunction:: dragon.vm.tensorrt.onnx.supports_device
.. raw:: html
<style>
h1:before {
content: "tensorrt.onnx.";
color: #103d3e;
}
</style>
docs/api/python/torch/onnx.rst
......@@ -6,7 +6,7 @@ vm.torch.onnx
Functions
---------
`export(...) <onnx/export.html>`_
: Export a model into ONNX format.
: Export the recorded graph to an onnx model.
.. toctree::
:hidden:
......
dragon/modules/python/module.cc
......@@ -274,7 +274,7 @@ PYBIND11_MODULE(libdragon_python, m) {
})
/*! \brief Load tensors and graph from a ONNX model */
.def("ImportONNXModel", [](Workspace* self, const string& model_path) {
.def("PrepareONNXModel", [](Workspace* self, const string& model_path) {
GraphDef init_graph, pred_graph;
onnx::ONNXBackend onnx_backend;
onnx_backend.Prepare(model_path, &init_graph, &pred_graph);
......
dragon/python/__init__.py
......@@ -29,6 +29,7 @@ from dragon._api import losses
from dragon._api import math
from dragon._api import metrics
from dragon._api import nn
from dragon._api import onnx
from dragon._api import optimizers
from dragon._api import random
from dragon._api import vision
......
dragon/python/_api/onnx/__init__.py 0 → 100644
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
from __future__ import absolute_import as _absolute_import
from __future__ import division as _division
from __future__ import print_function as _print_function
# Classes
from dragon.vm.onnx.core.backend.native import BackendRep
# Functions
from dragon.vm.onnx.core.backend.native import prepare as prepare_backend
from dragon.vm.onnx.core.backend.native import run_model
from dragon.vm.onnx.core.backend.native import supports_device
from dragon.vm.onnx.core.frontend.native import export
from dragon.vm.onnx.core.frontend.native import record
# Attributes
__all__ = [_s for _s in dir() if not _s.startswith('_')]
dragon/python/core/autograph/function_lib.py
......@@ -20,8 +20,8 @@ import os
from dragon.core.autograph import grad_maker
from dragon.core.autograph.op_def import OpDef
from dragon.core.autograph.op_def import OpInfo
from dragon.core.autograph.tensor import TensorRef
from dragon.core.framework import config
from dragon.core.framework import context
from dragon.core.framework import proto_util
from dragon.core.framework import types
from dragon.core.framework import workspace
......@@ -33,12 +33,14 @@ from dragon.core.util import nest
def add_device_option(graph_def):
"""Add the device option."""
cfg = config.config()
str2idx = {'cpu': 0, 'cuda': 1, 'cnml': 2}
dev_opt = dragon_pb2.DeviceOption()
dev_opt.device_type = str2idx[cfg.device_type]
dev_opt.device_id = cfg.device_index
dev_opt.random_seed = cfg.random_seed
graph_def.device_option.CopyFrom(dev_opt)
spec = context.get_device_spec()
graph_def.device_option.CopyFrom(
dragon_pb2.DeviceOption(
device_type={'cpu': 0,
'cuda': 1,
'cnml': 2}[spec.type],
device_id=spec.index,
random_seed=cfg.random_seed))
def add_grad_info(graph_def, targets):
......@@ -226,17 +228,13 @@ class Function(object):
f.write(str(graph_def))
logging.info('Export meta graph into: {}'.format(path))
def import_from(self, graph_def, explicit_inputs=False):
def import_from(self, graph_def):
"""Import a defined function from a graph def.
Set ``explicit_inputs`` to **True** to enforce feeding.
Parameters
----------
graph_def : GraphDef
The definition of graph.
explicit_inputs : bool
Whether to enforce feeding on executing.
Returns
-------
......@@ -244,8 +242,9 @@ class Function(object):
The self.
"""
self.outputs = [TensorRef(name) for name in graph_def.output]
self.inputs = [TensorRef(name).constant() for name in graph_def.input]
current_ws = workspace.get_workspace()
self.outputs = [current_ws.create_tensor(name) for name in graph_def.output]
self.inputs = [current_ws.create_tensor(name) for name in graph_def.input]
# Fill with all known graph elements.
add_device_option(graph_def)
......@@ -253,18 +252,14 @@ class Function(object):
add_phase(graph_def, self.outputs)
# Notify the backend to create and optimize.
current_ws = workspace.get_workspace()
graph_def.name = self.graph_def.name
self.graph_def = graph_def
self.graph_name = current_ws.create_graph(graph_def)
# Bind a callback to run this graph.
self.callback = lambda *args, **kwargs: \
current_ws.run_graph(
name=self.graph_name,
inputs_and_values=(self.inputs if explicit_inputs else [], args),
outputs=self.outputs,
**kwargs
)
name=self.graph_name, outputs=self.outputs, **kwargs)
return self
......
dragon/python/core/autograph/op_spec.py
......@@ -26,7 +26,6 @@ register = _GLOBAL_REGISTERED_SPECS.register
@register('Accuracy')
def accuracy_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype, outputs[0].shape = 'float32', []
return outputs
......@@ -97,7 +96,6 @@ def binary_shape_spec(inputs, outputs):
'Where',
])
def binary_math_spec(args, inputs, outputs):
_ = locals()
outputs = binary_shape_spec(inputs, outputs)
outputs[0].dtype = inputs[0].dtype
if inputs[0].dtype is None:
......@@ -114,7 +112,6 @@ def binary_math_spec(args, inputs, outputs):
'NotEqual',
])
def binary_compare_spec(args, inputs, outputs):
_ = locals()
outputs = binary_shape_spec(inputs, outputs)
outputs[0].dtype = 'bool'
return outputs
......@@ -262,7 +259,6 @@ def depth_to_space_spec(args, inputs, outputs):
@register('Dot')
def dot_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = inputs[0].dtype
try:
a_shape, b_shape = inputs[0].shape[:], inputs[1].shape[:]
......@@ -358,7 +354,6 @@ def expand_dims_spec(args, inputs, outputs):
'TruncatedNormal',
])
def fill_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = args['dtype']
try:
if 'dims' in args:
......@@ -476,7 +471,6 @@ def index_select_spec(args, inputs, outputs):
@register(['IsInf', 'IsNaN'])
def is_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = 'bool'
try:
outputs[0].shape = inputs[0].shape[:]
......@@ -487,7 +481,6 @@ def is_spec(args, inputs, outputs):
@register('LinSpace')
def linspace_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = args['dtype']
outputs[0].shape = args['dims']
return outputs
......@@ -495,7 +488,6 @@ def linspace_spec(args, inputs, outputs):
@register('MaskedSelect')
def masked_select_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = inputs[0].dtype
outputs[0].shape = (None,)
return outputs
......@@ -552,7 +544,6 @@ def multinomial_spec(args, inputs, outputs):
@register('NonZero')
def non_zero_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = 'int64'
try:
outputs[0].shape = (None, len(inputs[0].shape))
......@@ -592,7 +583,6 @@ def pad_spec(args, inputs, outputs):
@register('Permutation')
def permutation_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = args['dtype']
if len(inputs) == 1:
try:
......@@ -638,13 +628,11 @@ def pool_spec(args, inputs, outputs):
@register(['PythonPlugin', 'PythonPluginInfer'])
def python_spec(args, inputs, outputs):
_ = locals()
return outputs
@register('Range')
def range_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = args['dtype']
slice_args = args['slice']
if len(slice_args) == 2:
......@@ -802,7 +790,6 @@ def roi_pool_spec(args, inputs, outputs):
@register('Shape')
def shape_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = 'int64'
try:
outputs[0].shape = [len(inputs[0].shape)]
......@@ -861,7 +848,6 @@ def softmax_loss_spec(args, inputs, outputs):
@register('Sort')
def sort_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = inputs[0].dtype
outputs[1].dtype = 'int64'
try:
......@@ -1054,7 +1040,6 @@ def top_k_spec(args, inputs, outputs):
@register('Unchanged')
def unchanged_spec(args, inputs, outputs):
_ = locals()
outputs[0].dtype = inputs[0].dtype
try:
outputs[0].shape = inputs[0].shape[:]
......
dragon/python/core/util/deprecation.py
......@@ -80,7 +80,6 @@ def deprecated(date, instructions, warn_once=True):
def not_installed(package=''):
"""Return a dummy function for the package that is not installed."""
def dummy_fn(*args, **kwargs):
_ = locals()
raise ImportError('Package <%s> is required but not installed.' % package)
return dummy_fn
......
dragon/python/vm/onnx/__init__.py
......@@ -14,9 +14,6 @@ from __future__ import absolute_import as _absolute_import
from __future__ import division as _division
from __future__ import print_function as _print_function
# Classes
from dragon.vm.onnx.core.shell import Shell
# Attributes
from dragon.vm.onnx.core import nodes as _nodes
from dragon.vm.onnx.core import exporters as _import_exporters
__all__ = [_s for _s in dir() if not _s.startswith('_')]
dragon/python/vm/onnx/core/backend/native.py 0 → 100644
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
"""Native ONNX backend."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import os
import numpy
try:
import onnx
from onnx.backend.base import Backend
from onnx.backend.base import BackendRep as ONNXBackendRep
from onnx.backend.base import Device
from onnx.backend.base import DeviceType
from onnx.backend.base import namedtupledict
except ImportError:
from dragon.core.util import deprecation
onnx = deprecation.NotInstalled('onnx')
Backend = object
ONNXBackendRep = object
Device = deprecation.NotInstalled('onnx')
DeviceType = deprecation.NotInstalled('onnx')
from dragon.core.autograph import function_lib
from dragon.core.eager.tensor import EagerTensor
from dragon.core.framework import context
from dragon.core.framework import device_spec
from dragon.core.framework import workspace
from dragon.core.proto import dragon_pb2
from dragon.core.util import nest
class BackendRep(ONNXBackendRep):
"""ONNX-Dragon backend to execute repeatedly."""
def __init__(self, model, device, **kwargs):
"""Create a ``BackendRep``.
Parameters
----------
model : str
The path of onnx model file.
device : onnx.Device
The executing device.
"""
if not isinstance(device, Device):
device = Device(device)
graph_str = workspace.get_workspace().PrepareONNXModel(model)
graph_def = dragon_pb2.GraphDef()
graph_def.ParseFromString(graph_str)
if device.type == DeviceType.CPU:
device_type, device_index = 'cpu', 0
elif device.type == DeviceType.CUDA:
device_type, device_index = 'cuda', device.device_id
else:
raise ValueError('Unsupported device type: ' + device.type)
with context.device(device_type, device_index):
self._function = function_lib.Function(name='ONNXGraph') \
.import_from(graph_def)
self._input_dict = collections.OrderedDict(
[(impl.name, EagerTensor(impl=impl, device=device_spec.DeviceSpec(
device_type, device_index))) for impl in self._function.inputs])
self._output_dict = collections.OrderedDict(
[(impl.name, EagerTensor(impl=impl, device=device_spec.DeviceSpec(
device_type, device_index))) for impl in self._function.outputs])
def run(self, inputs, **kwargs):
"""Run the model.
Parameters
----------
inputs : Union[Sequence, Dict]
The input arrays.
Returns
-------
namedtuple
The model outputs.
"""
if isinstance(inputs, numpy.ndarray):
inputs = [inputs]
if isinstance(inputs, dict):
for name, value in inputs.items():
self._input_dict[name]._impl.FromNumpy(value)
elif nest.is_sequence(inputs):
for ref, value in zip(self._input_dict.values(), inputs):
ref._impl.FromNumpy(value)
else:
raise ValueError('Excepted sequence or dict inputs.')
self._function.callback(return_outputs=False)
named_outputs = namedtupledict('Outputs', list(self._output_dict.keys()))
return named_outputs(*(self._output_dict.values()))
class DragonBackend(Backend):
"""ONNX-Dragon backend."""
@classmethod
def prepare(cls, model, device='CPU:0', **kwargs):
"""Create a backend to execute repeatedly.
Parameters
----------
model : str
The path of onnx model file.
device : str, optional, default='CPU:0'
The executing device.
Returns
-------
dragon.onnx.BackendRep
The backend.
"""
if not os.path.exists(model):
raise ValueError('Model({}) is not existed.'.format(model))
return BackendRep(model, device, **kwargs)
@classmethod
def run_model(cls, model, inputs, device='CUDA:0', **kwargs):
"""Execute an onnx model once.
Parameters
----------
model : str
The path of onnx model file.
inputs : Union[Sequence, Dict]
The input arrays.
device : str, optional
The executing device.
Returns
-------
namedtuple
The model outputs.
"""
return cls.prepare(model, device, **kwargs).run(inputs)
@classmethod
def supports_device(cls, device_str):
"""Query if the given device is supported.
Parameters
----------
device_str : str
The device descriptor.
Returns
-------
bool
**True** if device is supported otherwise **False**.
"""
device = Device(device_str)
if device.type in (DeviceType.CPU, DeviceType.CUDA):
return True
return False
prepare = DragonBackend.prepare
run_model = DragonBackend.run_model
supports_device = DragonBackend.supports_device
tensorrt/python/core/backend.py → dragon/python/vm/onnx/core/backend/tensorrt.py
......@@ -7,11 +7,8 @@
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# Codes are based on:
#
# <https://github.com/onnx/onnx-tensorrt/blob/master/onnx_tensorrt/backend.py>
#
# ------------------------------------------------------------
"""TensorRT ONNX backend."""
from __future__ import absolute_import
from __future__ import division
......@@ -22,14 +19,14 @@ import numpy
try:
import onnx
from onnx.backend.base import Backend
from onnx.backend.base import BackendRep
from onnx.backend.base import BackendRep as ONNXBackendRep
from onnx.backend.base import Device
from onnx.backend.base import DeviceType
except ImportError:
from dragon.core.util import deprecation
onnx = deprecation.NotInstalled('onnx')
Backend = object
BackendRep = object
ONNXBackendRep = object
Device = deprecation.NotInstalled('onnx')
DeviceType = deprecation.NotInstalled('onnx')
......@@ -41,8 +38,8 @@ from dragon.vm.tensorrt.core.engine import trt
from dragon.vm.tensorrt.core.engine import TRT_LOGGER
class ONNXBackendRep(BackendRep):
"""Load and run onnx models."""
class BackendRep(ONNXBackendRep):
"""ONNX-TensorRT backend to execute repeatedly."""
def __init__(
self,
......@@ -53,7 +50,7 @@ class ONNXBackendRep(BackendRep):
optimization_profiles=None,
serialize_engine=False,
):
"""Create a ``ONNXBackendRep``.
"""Create a ``BackendRep``.
Parameters
----------
......@@ -71,6 +68,8 @@ class ONNXBackendRep(BackendRep):
Whether to serialize engine into a file.
"""
if not isinstance(device, Device):
device = Device(device)
self._set_device(device)
self._logger = TRT_LOGGER
self._builder = trt.Builder(self._logger)
......@@ -193,12 +192,12 @@ class ONNXBackendRep(BackendRep):
cuda.set_device(device.device_id)
class ONNXBackend(Backend):
class TensorRTBackend(Backend):
"""ONNX-TensorRT backend."""
@classmethod
def prepare(cls, model, device='CUDA:0', **kwargs):
"""Build a TensorRT engine from the onnx model.
"""Create a backend to execute repeatedly.
Parameters
----------
......@@ -209,17 +208,15 @@ class ONNXBackend(Backend):
Returns
-------
dragon.vm.tensorrt.ONNXBackendRep
The backend rep.
tensorrt.onnx.BackendRep
The backend.
"""
if not isinstance(device, Device):
device = Device(device)
return ONNXBackendRep(model, device, **kwargs)
return BackendRep(model, device, **kwargs)
@classmethod
def run_model(cls, model, inputs, device='CUDA:0', **kwargs):
"""Build and run a TensorRT engine from the onnx model.
"""Execute an onnx model once.
Parameters
----------
......@@ -240,7 +237,7 @@ class ONNXBackend(Backend):
@classmethod
def run_node(cls, node, inputs, device='CUDA:0', **kwargs):
"""Build and run a TensorRT engine from the onnx node.
"""Execute an onnx node once.
Parameters
----------
......@@ -248,7 +245,7 @@ class ONNXBackend(Backend):
The onnx node.
inputs : Union[Sequence, Dict]
The input arrays.
device : str, optional
device : str, optional, default='CUDA:0'
The executing device.
Returns
......@@ -257,7 +254,7 @@ class ONNXBackend(Backend):
The model outputs.
"""
super(ONNXBackend, cls).run_node(node, inputs, device)
super(TensorRTBackend, cls).run_node(node, inputs, device)
model = onnx_helper.make_model_from_node(node, inputs, use_weights=True)
try:
results = cls.prepare(model, device).run(inputs[:1])
......@@ -268,7 +265,7 @@ class ONNXBackend(Backend):
@classmethod
def supports_device(cls, device_str):
"""Query if given device is supported.
"""Query if the given device is supported.
Parameters
----------
......@@ -285,7 +282,7 @@ class ONNXBackend(Backend):
return device.type == DeviceType.CUDA
prepare = ONNXBackend.prepare
run_node = ONNXBackend.run_node
run_model = ONNXBackend.run_model
supports_device = ONNXBackend.supports_device
prepare = TensorRTBackend.prepare
run_node = TensorRTBackend.run_node
run_model = TensorRTBackend.run_model
supports_device = TensorRTBackend.supports_device
dragon/python/vm/onnx/core/nodes/__init__.py → dragon/python/vm/onnx/core/exporters/__init__.py
......@@ -13,8 +13,8 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.onnx.core.nodes import activation
from dragon.vm.onnx.core.nodes import array
from dragon.vm.onnx.core.nodes import normalization
from dragon.vm.onnx.core.nodes import math
from dragon.vm.onnx.core.nodes import vision
from dragon.vm.onnx.core.exporters import activation
from dragon.vm.onnx.core.exporters import array
from dragon.vm.onnx.core.exporters import normalization
from dragon.vm.onnx.core.exporters import math
from dragon.vm.onnx.core.exporters import vision
dragon/python/vm/onnx/core/nodes/activation.py → dragon/python/vm/onnx/core/exporters/activation.py
......@@ -13,26 +13,26 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.onnx.core import exporter
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
@exporter.register('Dropout')
def dropout_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Dropout')
def dropout_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
drop_ratio = 0.5 # The prob to set zeros randomly.
for arg in op_def.arg:
if arg.name == 'prob':
drop_ratio = arg.f
elif arg.name == 'prob_desc':
drop_ratio = helper.fetch_argument(op_def, arg, ws)
drop_ratio = helper.fetch_argument(op_def, arg, context.ws)
helper.add_attribute(node, 'ratio', drop_ratio)
return node, const_tensors
@exporter.register('HardSigmoid')
def hardsigmoid_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('HardSigmoid')
def hardsigmoid_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
alpha, beta = 0.2, 0.5
for arg in op_def.arg:
if arg.name == 'alpha':
......@@ -44,16 +44,16 @@ def hardsigmoid_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('PRelu')
def prelu_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
const_tensors = [helper.from_tensor(op_def.input[1], ws)]
@export_util.register('PRelu')
def prelu_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
const_tensors = [helper.from_tensor(op_def.input[1], context.ws)]
return node, const_tensors
@exporter.register('Relu')
def relu_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Relu')
def relu_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
if arg.name == 'alpha':
if arg.f > 0:
......@@ -62,9 +62,9 @@ def relu_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Selu')
def selu_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Selu')
def selu_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
alpha, gamma = 1.67326, 1.0507
for arg in op_def.arg:
if arg.name == 'alpha':
......@@ -76,10 +76,10 @@ def selu_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Softmax')
def softmax_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
ndim = len(shape_dict[op_def.input[0]])
@export_util.register('Softmax')
def softmax_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
ndim = len(context.blob_shapes[op_def.input[0]])
for arg in op_def.arg:
if arg.name == 'axis':
axis = arg.i + (ndim if arg.i < 0 else 0)
......
dragon/python/vm/onnx/core/nodes/array.py → dragon/python/vm/onnx/core/exporters/array.py
......@@ -20,13 +20,13 @@ except ImportError:
from dragon.core.util import deprecation
TensorProto = deprecation.not_installed('ONNX')
from dragon.vm.onnx.core import exporter
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
@exporter.register(['ArgMax', 'ArgMin'])
def arg_reduce_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register(['ArgMax', 'ArgMin'])
def arg_reduce_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
# ONNX requires indices only, remove the values.
indices = node.output[0]
node.ClearField('output')
......@@ -39,9 +39,9 @@ def arg_reduce_exporter(op_def, shape_dict, ws):
return node, None
@exporter.register('Cast')
def cast_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Cast')
def cast_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Cast'
if len(node.input) == 0:
raise ValueError('ONNX does not support in-place cast.')
......@@ -51,9 +51,9 @@ def cast_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('ChannelAffine')
def channel_affine_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('ChannelAffine')
def channel_affine_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'ATen' # Currently not supported in ai.onnx
helper.add_attribute(node, 'op_type', 'ChannelAffine')
for arg in op_def.arg:
......@@ -62,13 +62,13 @@ def channel_affine_exporter(op_def, shape_dict, ws):
elif arg.name == 'num_axes':
helper.add_attribute(node, 'num_axes', arg.i)
# Weights and biases
const_tensors = [helper.from_tensor(e, ws) for e in op_def.input[1:]]
const_tensors = [helper.from_tensor(e, context.ws) for e in op_def.input[1:]]
return node, const_tensors
@exporter.register('ChannelNormalize')
def channel_normalize_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('ChannelNormalize')
def channel_normalize_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'ATen' # Currently not supported in ai.onnx
helper.add_attribute(node, 'op_type', 'ChannelNormalize')
for arg in op_def.arg:
......@@ -83,27 +83,27 @@ def channel_normalize_exporter(op_def, shape_dict, ws):
elif arg.name == 'perm':
helper.add_attribute(node, 'perm', arg.ints)
elif arg.name == 'perm_desc':
values = helper.fetch_argument(op_def, arg, ws)
values = helper.fetch_argument(op_def, arg, context.ws)
helper.add_attribute(node, 'perm', values)
elif arg.name == 'perm_descs':
if len(arg.strings) > 0:
values = helper.fetch_arguments(op_def, arg, ws)
values = helper.fetch_arguments(op_def, arg, context.ws)
helper.add_attribute(node, 'perm', values)
return node, const_tensors
@exporter.register('Concat')
def concat_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Concat')
def concat_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
if arg.name == 'axis':
helper.add_attribute(node, 'axis', arg.i)
return node, const_tensors
@exporter.register('CumSum')
def cumulative_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('CumSum')
def cumulative_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
axis = 0
for arg in op_def.arg:
if arg.name == 'axis':
......@@ -114,27 +114,27 @@ def cumulative_exporter(op_def, shape_dict, ws):
helper.add_attribute(node, 'reverse', arg.i)
axis = helper.from_array(
numpy.array(axis, 'int64'),
op_def.input[0] + '/cumulative/axis',
context.unique_name(op_def.input[0] + '/cumulative/axis'),
)
node.input.extend([axis.name])
return node, [axis]
@exporter.register('Expand')
def expand_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
shape = list(shape_dict[op_def.output[0]])
@export_util.register('Expand')
def expand_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
shape = list(context.blob_shapes[op_def.output[0]])
shape = helper.from_array(
numpy.array(shape, 'int64'),
op_def.input[0] + '/expand/shape',
context.unique_name(op_def.input[0] + '/expand/shape'),
)
node.input.extend([shape.name])
return node, [shape]
@exporter.register('ExpandDims')
def expand_dims_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('ExpandDims')
def expand_dims_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Unsqueeze'
axes = None
for arg in op_def.arg:
......@@ -145,13 +145,13 @@ def expand_dims_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Eye')
def eye_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Eye')
def eye_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
if len(op_def.input) > 0:
node.op_type += 'Like'
else:
output_shape = list(shape_dict[op_def.output[0]])
output_shape = list(context.blob_shapes[op_def.output[0]])
helper.add_attribute(node, 'shape', output_shape)
for arg in op_def.arg:
if arg.name == 'k':
......@@ -161,9 +161,9 @@ def eye_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Flatten')
def flatten_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Flatten')
def flatten_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
if arg.name == 'axis':
helper.add_attribute(node, 'axis', arg.i)
......@@ -177,9 +177,9 @@ def flatten_exporter(op_def, shape_dict, ws):
return node, None
@exporter.register('IndexSelect')
def index_select_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('IndexSelect')
def index_select_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Gather'
for arg in op_def.arg:
if arg.name == 'axis':
......@@ -190,9 +190,9 @@ def index_select_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Multinomial')
def multinomial_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Multinomial')
def multinomial_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
helper.add_attribute(node, 'dtype', helper.tensor_type('int64'))
for arg in op_def.arg:
if arg.name == 'num_samples':
......@@ -200,12 +200,12 @@ def multinomial_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('OneHot')
def one_hot_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('OneHot')
def one_hot_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
helper.add_attribute(node, 'axis', -1)
depth, on_value, off_value = 1, 1, 0
dtype = ws.FetchTensor(node.output[0]).dtype
dtype = context.ws.FetchTensor(node.output[0]).dtype
for arg in op_def.arg:
if arg.name == 'depth':
depth = arg.i
......@@ -215,27 +215,27 @@ def one_hot_exporter(op_def, shape_dict, ws):
off_value = arg.i
depth = helper.from_array(
numpy.array(depth, 'int64'),
op_def.input[0] + '/one_hot/depth',
context.unique_name(op_def.input[0] + '/one_hot/depth'),
)
values = helper.from_array(
numpy.array([off_value, on_value], dtype),
op_def.input[0] + '/one_hot/values',
context.unique_name(op_def.input[0] + '/one_hot/values'),
)
const_tensors = [depth, values]
node.input.extend([depth.name, values.name])
return node, const_tensors
def pad_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
def pad_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
pads, value = [], 0
for arg in op_def.arg:
if arg.name == 'pads':
pads = [int(e) for e in arg.ints]
elif arg.name == 'pads_desc':
pads = helper.fetch_argument(op_def, arg, ws)
pads = helper.fetch_argument(op_def, arg, context.ws)
elif arg.name == 'pads_descs':
pads = helper.fetch_arguments(op_def, arg, ws)
pads = helper.fetch_arguments(op_def, arg, context.ws)
elif arg.name == 'mode':
helper.add_attribute(node, 'mode', arg.s.lower())
elif arg.name == 'value':
......@@ -243,44 +243,44 @@ def pad_exporter(op_def, shape_dict, ws):
return node, pads, value
@exporter.register('Pad-1')
def pad_exporter_v1(op_def, shape_dict, ws):
@export_util.register('Pad-1')
def pad_exporter_v1(op_def, context):
node, pads, value = pad_exporter(**locals())
helper.add_attribute(node, 'paddings', pads)
helper.add_attribute(node, 'value', value)
return node, []
@exporter.register('Pad-2')
def pad_exporter_v2(op_def, shape_dict, ws):
@export_util.register('Pad-2')
def pad_exporter_v2(op_def, context):
node, pads, value = pad_exporter(**locals())
helper.add_attribute(node, 'pads', pads)
helper.add_attribute(node, 'value', value)
return node, []
@exporter.register('Pad-11')
def pad_exporter_v11(op_def, shape_dict, ws):
@export_util.register('Pad-11')
def pad_exporter_v11(op_def, context):
node, pads, value = pad_exporter(**locals())
pads = helper.from_array(
numpy.array(pads, 'int64'),
op_def.input[0] + '/pad/pads',
context.unique_name(op_def.input[0] + '/pad/pads'),
)
value = helper.from_array(
numpy.array(value, 'float64'),
op_def.input[0] + '/pad/value',
context.unique_name(op_def.input[0] + '/pad/value'),
)
node.input.extend([pads.name, value.name])
return node, [pads, value]
@exporter.register('RandomNormal')
def random_normal_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('RandomNormal')
def random_normal_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
if len(op_def.input) > 0:
node.op_type += 'Like'
else:
output_shape = list(shape_dict[op_def.output[0]])
output_shape = list(context.blob_shapes[op_def.output[0]])
helper.add_attribute(node, 'shape', output_shape)
for arg in op_def.arg:
if arg.name == 'mean':
......@@ -292,13 +292,13 @@ def random_normal_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('RandomUniform')
def random_uniform_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('RandomUniform')
def random_uniform_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
if len(op_def.input) > 0:
node.op_type += 'Like'
else:
output_shape = list(shape_dict[op_def.output[0]])
output_shape = list(context.blob_shapes[op_def.output[0]])
helper.add_attribute(node, 'shape', output_shape)
for arg in op_def.arg:
if arg.name == 'low':
......@@ -310,10 +310,10 @@ def random_uniform_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register(['ReduceMax', 'ReduceMean', 'ReduceMin', 'ReduceSum'])
def reduce_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
axes = list(range(len(shape_dict[op_def.input[0]])))
@export_util.register(['ReduceMax', 'ReduceMean', 'ReduceMin', 'ReduceSum'])
def reduce_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
axes = list(range(len(context.blob_shapes[op_def.input[0]])))
for arg in op_def.arg:
if arg.name == 'axes':
axes = arg.ints
......@@ -323,44 +323,44 @@ def reduce_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Reshape')
def reshape_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
shape = dims = list(shape_dict[op_def.output[0]])
@export_util.register('Reshape')
def reshape_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
shape = dims = list(context.blob_shapes[op_def.output[0]])
for arg in op_def.arg:
if arg.name == 'dims':
dims = [int(e) for e in arg.ints]
elif arg.name == 'dims_desc':
dims = helper.fetch_argument(op_def, arg, ws)
dims = helper.fetch_argument(op_def, arg, context.ws)
elif arg.name == 'dims_descs':
dims = helper.fetch_arguments(op_def, arg, ws)
dims = helper.fetch_arguments(op_def, arg, context.ws)
for axis, dim in enumerate(dims):
shape[axis] = dim if dim <= 0 else shape[axis]
shape = helper.from_array(
numpy.array(shape, 'int64'),
op_def.input[0] + '/reshape/shape',
context.unique_name(op_def.input[0] + '/reshape/shape'),
)
node.input.extend([shape.name])
return node, [shape]
def slice_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
in_shape = shape_dict[op_def.input[0]]
def slice_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
in_shape = context.blob_shapes[op_def.input[0]]
starts, sizes, ends = [], [], []
for arg in op_def.arg:
if arg.name == 'starts':
starts = [int(e) for e in arg.ints]
elif arg.name == 'starts_desc':
starts = helper.fetch_argument(op_def, arg, ws)
starts = helper.fetch_argument(op_def, arg, context.ws)
elif arg.name == 'starts_descs':
starts = helper.fetch_arguments(op_def, arg, ws)
starts = helper.fetch_arguments(op_def, arg, context.ws)
elif arg.name == 'sizes':
sizes = [int(e) for e in arg.ints]
elif arg.name == 'sizes_desc':
sizes = helper.fetch_argument(op_def, arg, ws)
sizes = helper.fetch_argument(op_def, arg, context.ws)
elif arg.name == 'sizes_descs':
sizes = helper.fetch_arguments(op_def, arg, ws)
sizes = helper.fetch_arguments(op_def, arg, context.ws)
for i, size in enumerate(sizes):
if size == -1:
ends.append(in_shape[i])
......@@ -371,8 +371,8 @@ def slice_exporter(op_def, shape_dict, ws):
return node, starts, ends
@exporter.register('Slice-1')
def slice_exporter_v1(op_def, shape_dict, ws):
@export_util.register('Slice-1')
def slice_exporter_v1(op_def, context):
node, starts, ends = slice_exporter(**locals())
helper.add_attribute(node, 'axes', numpy.arange(len(starts)))
helper.add_attribute(node, 'ends', ends)
......@@ -380,40 +380,40 @@ def slice_exporter_v1(op_def, shape_dict, ws):
return node, []
@exporter.register('Slice-10')
def slice_exporter_v10(op_def, shape_dict, ws):
@export_util.register('Slice-10')
def slice_exporter_v10(op_def, context):
node, starts, ends = slice_exporter(**locals())
axes = helper.from_array(
numpy.arange(len(starts), dtype='int64'),
op_def.input[0] + '/slice/axes',
context.unique_name(op_def.input[0] + '/slice/axes'),
)
starts = helper.from_array(
numpy.array(starts, 'int64'),
op_def.input[0] + '/slice/starts',
context.unique_name(op_def.input[0] + '/slice/starts'),
)
ends = helper.from_array(
numpy.array(ends, 'int64'),
op_def.input[0] + '/slice/ends',
context.unique_name(op_def.input[0] + '/slice/ends'),
)
node.input.extend([starts.name, ends.name, axes.name])
return node, [starts, ends, axes]
@exporter.register('Split')
def split_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Split')
def split_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
axis = 0
for arg in op_def.arg:
if arg.name == 'axis':
axis = arg.i
size_splits = [shape_dict[e][axis] for e in op_def.output]
size_splits = [context.blob_shapes[e][axis] for e in op_def.output]
helper.add_attribute(node, 'split', size_splits)
return node, const_tensors
@exporter.register('Squeeze')
def squeeze_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Squeeze')
def squeeze_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
axes = None
for arg in op_def.arg:
if arg.name == 'axes':
......@@ -423,43 +423,43 @@ def squeeze_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Tile')
def tile_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Tile')
def tile_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
repeats = []
for arg in op_def.arg:
if arg.name == 'repeats':
repeats = [e for e in arg.ints]
elif arg.name == 'repeats_desc':
repeats = helper.fetch_argument(op_def, arg, ws)
repeats = helper.fetch_argument(op_def, arg, context.ws)
elif arg.name == 'repeats_descs':
repeats = helper.fetch_arguments(op_def, arg, ws)
repeats = helper.fetch_arguments(op_def, arg, context.ws)
repeats = helper.from_array(
numpy.array(repeats, 'int64'),
op_def.input[0] + '/tile/repeats',
context.unique_name(op_def.input[0] + '/tile/repeats'),
)
node.input.extend([repeats.name])
return node, [repeats]
@exporter.register('Transpose')
def transpose_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Transpose')
def transpose_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
if arg.name == 'perm':
helper.add_attribute(node, 'perm', arg.ints)
elif arg.name == 'perm_desc':
values = helper.fetch_argument(op_def, arg, ws)
values = helper.fetch_argument(op_def, arg, context.ws)
helper.add_attribute(node, 'perm', values)
elif arg.name == 'perm_descs':
if len(arg.strings) > 0:
values = helper.fetch_arguments(op_def, arg, ws)
values = helper.fetch_arguments(op_def, arg, context.ws)
helper.add_attribute(node, 'perm', values)
return node, None
def top_k_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
def top_k_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
k, axis, largest, sorted = 1, -1, True, True
for arg in op_def.arg:
if arg.name == 'k':
......@@ -473,9 +473,9 @@ def top_k_exporter(op_def, shape_dict, ws):
return node, (k, axis, largest, sorted)
@exporter.register('TopK-1')
def top_k_exporter_v1(op_def, shape_dict, ws):
node, (k, axis, largest, sorted) = top_k_exporter(op_def, shape_dict, ws)
@export_util.register('TopK-1')
def top_k_exporter_v1(op_def, context):
node, (k, axis, largest, sorted) = top_k_exporter(**locals())
if largest == 0:
raise ValueError('TopK-1 does not support smallest mode.')
helper.add_attribute(node, 'axis', axis)
......@@ -483,37 +483,37 @@ def top_k_exporter_v1(op_def, shape_dict, ws):
return node, None
@exporter.register('TopK-10')
def top_k_exporter_v10(op_def, shape_dict, ws):
node, (k, axis, largest, sorted) = top_k_exporter(op_def, shape_dict, ws)
@export_util.register('TopK-10')
def top_k_exporter_v10(op_def, context):
node, (k, axis, largest, sorted) = top_k_exporter(**locals())
if largest == 0:
raise ValueError('TopK-10 does not support smallest mode.')
helper.add_attribute(node, 'axis', axis)
k = helper.from_array(
numpy.array([k], 'int64'),
op_def.input[0] + '/top_k/k',
context.unique_name(op_def.input[0] + '/top_k/k'),
)
node.input.extend([k.name])
return node, [k]
@exporter.register('TopK-11')
def top_k_exporter_v11(op_def, shape_dict, ws):
node, (k, axis, largest, sorted) = top_k_exporter(op_def, shape_dict, ws)
@export_util.register('TopK-11')
def top_k_exporter_v11(op_def, context):
node, (k, axis, largest, sorted) = top_k_exporter(**locals())
helper.add_attribute(node, 'axis', axis)
helper.add_attribute(node, 'largest', largest)
helper.add_attribute(node, 'sorted', sorted)
k = helper.from_array(
numpy.array([k], 'int64'),
op_def.input[0] + '/top_k/k',
context.unique_name(op_def.input[0] + '/top_k/k'),
)
node.input.extend([k.name])
return node, [k]
@exporter.register('Unique')
def unique_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Unique')
def unique_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
helper.add_attribute(node, 'sorted', 1)
return_inverse = return_counts = 0
for arg in op_def.arg:
......
dragon/python/vm/onnx/core/nodes/math.py → dragon/python/vm/onnx/core/exporters/math.py
......@@ -15,35 +15,35 @@ from __future__ import print_function
import numpy
from dragon.vm.onnx.core import exporter
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
@exporter.register('Add')
def add_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], ws).dtype)
@export_util.register('Add')
def add_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], context.ws).dtype)
node.op_type = 'Or' if dtype == 'bool' else 'Add'
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Div')
def div_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Div')
def div_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Clip')
def clip_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Clip')
def clip_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
if arg.name == 'low':
helper.add_attribute(node, 'min', arg.f)
......@@ -52,11 +52,11 @@ def clip_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Clip-11')
def clip_exporter_v11(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Clip-11')
def clip_exporter_v11(op_def, context):
node, const_tensors = export_util.translate(**locals())
min_value, max_value, const_tensors = None, None, []
dtype = ws.FetchTensor(op_def.output[0]).dtype
dtype = context.ws.FetchTensor(op_def.output[0]).dtype
for arg in op_def.arg:
if arg.name == 'low':
min_value = arg.f
......@@ -65,7 +65,7 @@ def clip_exporter_v11(op_def, shape_dict, ws):
if min_value is not None:
const_tensors.append(helper.from_array(
numpy.array(min_value, dtype),
op_def.input[0] + '/clip/min_value',
context.unique_name(op_def.input[0] + '/clip/min_value'),
))
node.input.extend([const_tensors[-1].name])
else:
......@@ -73,7 +73,7 @@ def clip_exporter_v11(op_def, shape_dict, ws):
if max_value is not None:
const_tensors.append(helper.from_array(
numpy.array(max_value, dtype),
op_def.input[0] + '/clip/max_value',
context.unique_name(op_def.input[0] + '/clip/max_value'),
))
node.input.extend([const_tensors[-1].name])
else:
......@@ -81,9 +81,9 @@ def clip_exporter_v11(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('FullyConnected-7')
def fully_connected_exporter_v7(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('FullyConnected-7')
def fully_connected_exporter_v7(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Gemm'
helper.add_attribute(node, 'alpha', 1.)
helper.add_attribute(node, 'beta', 1.)
......@@ -91,27 +91,28 @@ def fully_connected_exporter_v7(op_def, shape_dict, ws):
if arg.name == 'transW':
helper.add_attribute(node, 'transB', arg.i)
# Weights and biases
const_tensors = [helper.from_tensor(e, ws) for e in op_def.input[1:]]
const_tensors = [helper.from_tensor(name, context.ws)
for name in op_def.input[1:]]
return node, const_tensors
@exporter.register('FullyConnected')
def fully_connected_exporter(op_def, shape_dict, ws):
node, const_tensors = fully_connected_exporter_v7(op_def, shape_dict, ws)
@export_util.register('FullyConnected')
def fully_connected_exporter(op_def, context):
node, const_tensors = fully_connected_exporter_v7(op_def, context)
helper.add_attribute(node, 'broadcast', 1) # Removed since opset 7
return node, const_tensors
@exporter.register('Invert')
def invert_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Invert')
def invert_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Not'
return node, const_tensors
@exporter.register('Matmul')
def matmul_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Matmul')
def matmul_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'MatMul'
for arg in op_def.arg:
if arg.name == 'transA':
......@@ -123,57 +124,57 @@ def matmul_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Maximum')
def maximum_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Maximum')
def maximum_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Max' # Eltwise, Broadcast
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Minimum')
def minimum_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Minimum')
def minimum_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Min' # Eltwise, Broadcast
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Mul')
def mul_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], ws).dtype)
@export_util.register('Mul')
def mul_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], context.ws).dtype)
node.op_type = 'And' if dtype == 'bool' else 'Mul'
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Pow')
def pow_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Pow')
def pow_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
@exporter.register('Sub')
def sub_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], ws).dtype)
@export_util.register('Sub')
def sub_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
dtype = str(helper.fetch_tensor(op_def.output[0], context.ws).dtype)
node.op_type = 'Xor' if dtype == 'bool' else 'Sub'
const_tensors = [] # Global scalars
for e in op_def.input:
if e.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(e, ws))
for name in op_def.input:
if name.startswith('/share/scalar/'):
const_tensors.append(helper.from_tensor(name, context.ws))
return node, const_tensors
dragon/python/vm/onnx/core/nodes/normalization.py → dragon/python/vm/onnx/core/exporters/normalization.py
......@@ -13,13 +13,13 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from dragon.vm.onnx.core import exporter
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
@exporter.register('BatchNorm')
def batch_norm_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('BatchNorm')
def batch_norm_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'BatchNormalization'
for arg in op_def.arg:
if arg.name == 'epsilon':
......@@ -27,13 +27,13 @@ def batch_norm_exporter(op_def, shape_dict, ws):
elif arg.name == 'momentum':
helper.add_attribute(node, 'momentum', arg.f)
# Weight, bias, running mean and running variance
const_tensors = [helper.from_tensor(e, ws) for e in op_def.input[1:]]
const_tensors = [helper.from_tensor(e, context.ws) for e in op_def.input[1:]]
return node, const_tensors
@exporter.register('GroupNorm')
def group_norm_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('GroupNorm')
def group_norm_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'ATen' # Currently not supported in ai.onnx
for arg in op_def.arg:
if arg.name == 'epsilon':
......@@ -46,13 +46,13 @@ def group_norm_exporter(op_def, shape_dict, ws):
helper.add_attribute(node, 'op_type', 'GroupNorm')
helper.add_attribute(node, 'group', arg.i)
# Weight and bias
const_tensors = [helper.from_tensor(e, ws) for e in op_def.input[1:]]
const_tensors = [helper.from_tensor(e, context.ws) for e in op_def.input[1:]]
return node, const_tensors
@exporter.register('LpNormalize')
def lp_normalize_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('LpNormalize')
def lp_normalize_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'LpNormalization'
for arg in op_def.arg:
if arg.name == 'axis':
......
dragon/python/vm/onnx/core/exporter.py → dragon/python/vm/onnx/core/exporters/utils.py
......@@ -21,13 +21,39 @@ _GLOBAL_REGISTERED_EXPORTERS = _Registry('exporters')
register = _GLOBAL_REGISTERED_EXPORTERS.register
def translate(op_def, *args, **kwargs):
class TranslatorContext(object):
"""Context to pass translator resources."""
def __init__(
self,
workspace,
blob_names,
blob_shapes,
blob_versions,
opset_version,
):
self.ws = workspace
self.blob_names = blob_names
self.blob_shapes = blob_shapes
self.blob_versions = blob_versions
self.opset_version = opset_version
def unique_name(self, name):
self.blob_versions[name] += 1
if self.blob_versions[name] > 1:
return name + '_%d' % (self.blob_versions[name] - 1)
return name
def translate(op_def, context):
"""Translate the OpDef to a NodeProto.
Parameters
----------
op_def : OperatorDef
The definition of a operator.
context : TranslatorContext
The context of translator.
Returns
-------
......@@ -37,9 +63,8 @@ def translate(op_def, *args, **kwargs):
The constant tensors.
"""
_ = locals()
node = helper.make_node(
op_type=kwargs.get('op_type', op_def.type),
op_type=op_def.type,
inputs=op_def.input,
outputs=op_def.output,
name=op_def.name if op_def.name != '' else None
......@@ -61,7 +86,6 @@ def registered_exporters():
@register(['PythonPlugin', 'PythonPluginInfer'])
def _python(op_def, shape_dict, ws):
def python_exporter(op_def, context):
"""Export the python operators."""
_ = locals()
return None, None
dragon/python/vm/onnx/core/nodes/vision.py → dragon/python/vm/onnx/core/exporters/vision.py
......@@ -16,44 +16,40 @@ from __future__ import print_function
import copy
import numpy
from dragon.vm.onnx.core import exporter
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
@exporter.register([
@export_util.register([
'Conv2d',
'ConvTranspose2d',
'DepthwiseConv2d',
])
def convolution(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
def convolution(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'ConvTranspose' if 'Transpose' in op_def.type else 'Conv'
if 'Depthwise' in op_def.type:
input_shape = shape_dict[op_def.input[0]]
input_shape = context.blob_shapes[op_def.input[0]]
helper.add_attribute(node, 'group', input_shape[1])
rank = len(shape_dict[op_def.input[0]]) - 2
rank = len(context.blob_shapes[op_def.input[0]]) - 2
for arg in op_def.arg:
_assert_data_format(arg)
if arg.name == 'kernel_shape':
helper.add_attribute(
node, 'kernel_shape',
_normalize_tuple(arg.ints, rank)
)
_normalize_tuple(arg.ints, rank))
elif arg.name == 'dilations':
helper.add_attribute(
node, 'dilations',
_normalize_tuple(arg.ints, rank)
)
_normalize_tuple(arg.ints, rank))
elif arg.name == 'strides':
helper.add_attribute(
node, 'strides',
_normalize_tuple(arg.ints, rank)
)
_normalize_tuple(arg.ints, rank))
elif arg.name == 'pads':
helper.add_attribute(
node, 'pads',
_normalize_pads(arg.ints, rank)
)
_normalize_pads(arg.ints, rank))
elif arg.name == 'padding' and arg.s != b'VALID':
helper.add_attribute(node, 'auto_pad', arg.s)
elif arg.name == 'group':
......@@ -63,13 +59,13 @@ def convolution(op_def, shape_dict, ws):
elif arg.name == 'output_padding':
helper.add_attribute(node, 'output_padding', arg.ints)
# Weights and biases
const_tensors = [helper.from_tensor(e, ws) for e in op_def.input[1:]]
const_tensors = [helper.from_tensor(e, context.ws) for e in op_def.input[1:]]
return node, const_tensors
@exporter.register(['DepthToSpace', 'SpaceToDepth'])
def depth_space_exporter(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register(['DepthToSpace', 'SpaceToDepth'])
def depth_space_exporter(op_def, context):
node, const_tensors = export_util.translate(**locals())
for arg in op_def.arg:
_assert_data_format(arg)
if arg.name == 'block_size':
......@@ -77,28 +73,25 @@ def depth_space_exporter(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Pool2d')
def pool(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
rank = len(shape_dict[op_def.input[0]]) - 2
@export_util.register('Pool2d')
def pool(op_def, context):
node, const_tensors = export_util.translate(**locals())
rank = len(context.blob_shapes[op_def.input[0]]) - 2
global_pooling, node_copy = 0, copy.deepcopy(node)
for arg in op_def.arg:
_assert_data_format(arg)
if arg.name == 'kernel_shape':
helper.add_attribute(
node, 'kernel_shape',
_normalize_tuple(arg.ints, rank)
)
_normalize_tuple(arg.ints, rank))
elif arg.name == 'strides':
helper.add_attribute(
node, 'strides',
_normalize_tuple(arg.ints, rank)
)
_normalize_tuple(arg.ints, rank))
elif arg.name == 'pads':
helper.add_attribute(
node, 'pads',
_normalize_pads(arg.ints, rank)
)
_normalize_pads(arg.ints, rank))
elif arg.name == 'padding' and arg.s != b'VALID':
helper.add_attribute(node, 'auto_pad', arg.s)
elif arg.name == 'mode':
......@@ -117,59 +110,53 @@ def pool(op_def, shape_dict, ws):
return node, const_tensors
@exporter.register('Resize-1')
def resize_v1(op_def, shape_dict, ws):
_ = locals()
@export_util.register('Resize-1')
def resize_v1(op_def, context):
raise RuntimeError('<Upsample> requires opset version >= 7.')
@exporter.register('Resize-7')
def resize_v7(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('Resize-7')
def resize_v7(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'Upsample'
input_shape = shape_dict[op_def.input[0]]
output_shape = shape_dict[op_def.output[0]]
input_shape = context.blob_shapes[op_def.input[0]]
output_shape = context.blob_shapes[op_def.output[0]]
for arg in op_def.arg:
if arg.name == 'mode':
helper.add_attribute(node, 'mode', arg.s.lower())
helper.add_attribute(
node, 'scales', [
float(output_shape[i]) / input_shape[i]
for i in range(len(input_shape))
]
)
node, 'scales', [float(output_shape[i]) / input_shape[i]
for i in range(len(input_shape))])
return node, const_tensors
@exporter.register('Resize-9')
def resize_v9(op_def, shape_dict, ws):
node, const_tensors = resize_v10(op_def, shape_dict, ws)
@export_util.register('Resize-9')
def resize_v9(op_def, context):
node, const_tensors = resize_v10(**locals())
node.op_type = 'Upsample'
return node, const_tensors
@exporter.register('Resize-10')
def resize_v10(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
input_shape = shape_dict[op_def.input[0]]
output_shape = shape_dict[op_def.output[0]]
@export_util.register('Resize-10')
def resize_v10(op_def, context):
node, const_tensors = export_util.translate(**locals())
input_shape = context.blob_shapes[op_def.input[0]]
output_shape = context.blob_shapes[op_def.output[0]]
for arg in op_def.arg:
if arg.name == 'mode':
helper.add_attribute(node, 'mode', arg.s.lower())
scales = helper.from_array(
numpy.array([
float(output_shape[i]) / input_shape[i]
for i in range(len(input_shape))
], 'float32'),
op_def.input[0] + '/resize/scales',
numpy.array([float(output_shape[i]) / input_shape[i]
for i in range(len(input_shape))], 'float32'),
context.unique_name(op_def.input[0] + '/resize/scales'),
)
node.input.extend([scales.name])
return node, [scales]
@exporter.register('Resize-11')
def resize_v11(op_def, shape_dict, ws):
node, const_tensors = resize_v10(op_def, shape_dict, ws)
@export_util.register('Resize-11')
def resize_v11(op_def, context):
node, const_tensors = resize_v10(**locals())
coord_mode = 'half_pixel'
for arg in op_def.arg:
if arg.name == 'mode':
......@@ -179,22 +166,22 @@ def resize_v11(op_def, shape_dict, ws):
if arg.i > 0:
coord_mode = 'align_corners'
helper.add_attribute(node, 'coordinate_transformation_mode', coord_mode)
rank = len(shape_dict[op_def.input[0]])
rank = len(context.blob_shapes[op_def.input[0]])
roi = helper.from_array(
numpy.array(([0] * rank + [1] * rank), 'float32'),
op_def.input[0] + '/resize/roi',
context.unique_name(op_def.input[0] + '/resize/roi'),
)
node.input[:] = [node.input[0], roi.name, node.input[1]]
return node, const_tensors + [roi]
@exporter.register('RoiAlign')
def roi_align(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('RoiAlign')
def roi_align(op_def, context):
node, const_tensors = export_util.translate(**locals())
# Make a dummy "batch_indices".
batch_indices = helper.from_array(
numpy.array([1], 'int64'),
op_def.input[0] + '/roi_align/batch_indices',
context.unique_name(op_def.input[0] + '/roi_align/batch_indices'),
)
node.input.extend([batch_indices.name])
for arg in op_def.arg:
......@@ -209,9 +196,9 @@ def roi_align(op_def, shape_dict, ws):
return node, [batch_indices]
@exporter.register('RoiPool')
def roi_pool(op_def, shape_dict, ws):
node, const_tensors = exporter.translate(**locals())
@export_util.register('RoiPool')
def roi_pool(op_def, context):
node, const_tensors = export_util.translate(**locals())
node.op_type = 'MaxRoiPool'
pooled_shape = [None, None]
for arg in op_def.arg:
......
tensorrt/python/_api/backend/__init__.py → dragon/python/vm/onnx/core/frontend/__init__.py
......@@ -8,14 +8,3 @@
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
from __future__ import absolute_import as _absolute_import
from __future__ import division as _division
from __future__ import print_function as _print_function
from dragon.vm.tensorrt.core.backend import prepare
from dragon.vm.tensorrt.core.backend import run_model
from dragon.vm.tensorrt.core.backend import run_node
from dragon.vm.tensorrt.core.backend import supports_device
__all__ = [_s for _s in dir() if not _s.startswith('_')]
dragon/python/vm/onnx/core/frontend.py → dragon/python/vm/onnx/core/frontend/native.py
......@@ -7,11 +7,8 @@
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# Codes are based on:
#
# <https://github.com/pytorch/pytorch/blob/master/caffe2/python/onnx/frontend.py>
#
# ------------------------------------------------------------
"""Native ONNX frontend."""
from __future__ import absolute_import
from __future__ import division
......@@ -20,14 +17,22 @@ from __future__ import print_function
import collections
import itertools
import numpy
try:
import onnx
except ImportError:
onnx = None
from dragon.core.autograph import function_lib
from dragon.core.eager import context as eager_context
from dragon.core.eager import backprop
from dragon.core.framework import types
from dragon.core.framework import workspace as workspace_util
from dragon.core.proto import dragon_pb2
from dragon.core.util import nest
from dragon.vm.onnx.core import exporter
from dragon.core.util import serialization
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.exporters import utils as export_util
class DragonFrontend(object):
......@@ -54,7 +59,7 @@ class DragonFrontend(object):
constants=None,
value_info=None,
opset_version=None,
ws=None,
workspace=None,
verbose=True,
):
input_names = [] if input_names is None else input_names
......@@ -79,12 +84,12 @@ class DragonFrontend(object):
blob_aliases = {}
for i, alias in enumerate(output_names):
blob_aliases[graph_def.output[i]] = alias
ws.RegisterAlias(graph_def.output[i], alias)
workspace.RegisterAlias(graph_def.output[i], alias)
if graph_def.output[i] in value_info:
value_info[alias] = value_info[graph_def.output[i]]
for i, alias in enumerate(input_names):
blob_aliases[graph_def.input[i]] = alias
ws.RegisterAlias(graph_def.input[i], alias)
workspace.RegisterAlias(graph_def.input[i], alias)
if graph_def.input[i] in value_info:
value_info[alias] = value_info[graph_def.input[i]]
......@@ -100,62 +105,65 @@ class DragonFrontend(object):
value_info[k] = (value_info[k][0], v)
# Prepare to make the graph.
onnx_graph = onnx.GraphProto(
name=graph_def.name
onnx_graph = onnx.GraphProto(name=graph_def.name
if len(graph_def.name) > 0
else 'onnx-model'
else 'onnx-model')
blob_shapes, blob_names = {}, {}
blob_versions = collections.defaultdict(
int, **dict((blob_aliases.get(k, k), 1)
for k in helper.collect_inputs(graph_def)))
initializers, seen_initializers = [], set()
# Build translator context.
context = export_util.TranslatorContext(
workspace=workspace,
blob_names=blob_names,
blob_shapes=blob_shapes,
blob_versions=blob_versions,
opset_version=opset_version,
)
graph_inputs = helper.collect_inputs(graph_def)
# Add nodes.
shapes, blob_names, initializers = {}, {}, []
blob_versions = collections.defaultdict(
int, **dict((blob_aliases.get(k, k), 1) for k in graph_inputs))
for op in graph_def.op:
# Get the shape of inputs and outputs.
for name in itertools.chain(op.input, op.output):
impl = ws.GetTensor(name)
impl = workspace.GetTensor(name)
if impl is not None:
shapes[name] = impl.dims
blob_shapes[name] = impl.dims
else:
shapes[name] = value_info[name][1]
blob_shapes[name] = value_info[name][1]
# Translate definition.
nodes, const_tensors = cls._translate(op, opset_version, shapes, ws)
nodes, const_tensors = cls._make_node(op, context)
# Rewritten for names.
for node in nodes:
node.input[:] = [blob_aliases.get(e, e) for e in node.input]
node.output[:] = [blob_aliases.get(e, e) for e in node.output]
node, blob_names, blob_versions = \
cls._ssa_rewrite(node, blob_names, blob_versions)
cls._rewrite_for_ssa(node, context)
# Directly convert outputs as const tensors if necessary.
# Convert constant outputs if necessary.
if None in nodes:
const_tensors = [helper.from_tensor(name, ws) for name in op.output]
const_tensors = [helper.from_tensor(name, workspace)
for name in op.output]
else:
onnx_graph.node.extend(nodes)
# Merge constant tensors.
if const_tensors is not None:
value_info = {
**value_info,
**dict((
e.name, (e.data_type, e.dims)
) for e in const_tensors)
}
initializers.extend(const_tensors)
value_info = {**value_info,
**dict((e.name, (e.data_type, e.dims))
for e in const_tensors)}
for tensor in const_tensors:
if tensor.name not in seen_initializers:
initializers.append(tensor)
seen_initializers.add(tensor.name)
# Add constants.
if constants is not None:
for k, v in constants.items():
initializers.append(helper.from_array(v, name=k))
# Add initializers.
onnx_graph.initializer.extend(initializers)
# Add inputs.
for name in helper.collect_inputs(onnx_graph):
try:
......@@ -165,18 +173,33 @@ class DragonFrontend(object):
elem_type=value_info[name][0],
shape=value_info[name][1])])
except KeyError:
impl = workspace.GetTensor(name)
if impl is not None:
initializer = helper.from_tensor(name, workspace)
onnx_graph.input.extend([
helper.make_tensor_value_info(
name=name,
elem_type=initializer.data_type,
shape=initializer.dims)])
if name not in seen_initializers:
initializers.append(initializer)
seen_initializers.add(initializer.name)
else:
raise ValueError(
'Info of tensor `{}` is missing, '
'specify it in <value_info>.'.format(name))
# Add initializers.
onnx_graph.initializer.extend(initializers)
# Add outputs.
onnx_graph.output.extend(
helper.make_tensor_value_info(
name=blob_names.get(name, name),
elem_type=value_info[name][0],
shape=value_info[name][1],
) for name in [blob_aliases.get(e, e) for e in set(graph_def.output)]
)
name=blob_names.get(name_v2, name_v2),
elem_type=value_info[name_v2][0],
shape=value_info[name_v2][1])
for name_v2 in [blob_aliases.get(name, name)
for name in set(graph_def.output)])
if verbose:
print(helper.printable_graph(onnx_graph))
......@@ -239,30 +262,31 @@ class DragonFrontend(object):
return opset_version
@staticmethod
def _ssa_rewrite(op_def, blob_names, blob_versions):
def _rewrite_for_ssa(op_def, context):
"""Rewrite a OpDef to satisfy the SSA (Static Single Assignment)."""
blob_names = context.blob_names
blob_versions = context.blob_versions
inputs, outputs = [], []
for e in op_def.input:
inputs.append(blob_names[e] if e in blob_names else e)
for e in op_def.output:
outputs.append(e + '_%d' % blob_versions[e]
if blob_versions[e] > 0 else e)
if e != '':
blob_versions[e] += 1
blob_names[e] = outputs[-1]
for name in op_def.input:
inputs.append(blob_names[name] if name in blob_names else name)
for name in op_def.output:
outputs.append(name + '_%d' % blob_versions[name]
if blob_versions[name] > 0 else name)
if name != '':
blob_versions[name] += 1
blob_names[name] = outputs[-1]
op_def.ClearField('input')
op_def.ClearField('output')
op_def.input.extend(inputs)
op_def.output.extend(outputs)
return op_def, blob_names, blob_versions
@classmethod
def _translate(cls, op_def, opset_version, shape_dict, ws):
def _make_node(cls, op_def, context):
"""Return a NodeProto from the OpDef."""
translate_fn = None
getter = exporter._GLOBAL_REGISTERED_EXPORTERS.try_get
getter = export_util._GLOBAL_REGISTERED_EXPORTERS.try_get
# Select the last versioned exporter if necessary.
for i in range(opset_version, 0, -1):
for i in range(context.opset_version, 0, -1):
versioned_op_type = op_def.type + '-%d' % i
if getter(versioned_op_type) is not None:
translate_fn = getter(versioned_op_type)
......@@ -273,10 +297,153 @@ class DragonFrontend(object):
translate_fn = getter(op_def.type)
else:
# Fallback to the generic exporter.
translate_fn = exporter.translate
nodes, const_tensors = translate_fn(op_def, shape_dict, ws)
translate_fn = export_util.translate
nodes, const_tensors = translate_fn(op_def, context)
return nest.flatten(nodes), const_tensors
def record():
"""Context-manger to record the graph.
Examples:
```python
with dragon.onnx.record():
...
```
See Also
--------
`dragon.onnx.export(...)`_
"""
tape = backprop.Tape()
tape.retain_graph = True
return backprop._GLOBAL_TAPE_STACK.get_controller(tape)
def export(
inputs,
outputs,
f,
input_names=None,
output_names=None,
input_shapes=None,
opset_version=None,
verbose=False,
enable_onnx_checker=True,
):
"""Export the recorded graph to an onnx model.
Enter into the record mode to export operators into an onnx model:
```python
x = dragon.constant([1, 2, 3])
with dragon.onnx.record():
y = x * x
dragon.onnx.export(inputs=[x], outputs=[y], f='model.onnx')
```
Parameters
----------
inputs : Union[Sequence, Dict]
The model inputs.
outputs : Union[Sequence, Dict]
The model outputs.
f : str
The filename for exporting model.
input_names : Sequence[str], optional
The name to the inputs.
output_names : Sequence[str], optional
The name to the outputs.
input_shapes : Union[Sequence, Dict], optional
The optional rewritten for input shapes.
opset_version : int, optional
The version of operator set.
verbose : bool, optional, default=False
Whether to print the debug string of graph.
enable_onnx_checker : bool, optional, default=True
Whether to check if model is valid.
"""
# Process the inputs.
if isinstance(inputs, dict):
if input_names is not None:
raise ValueError(
'Excepted the input names from <inputs>.\n'
'You should set the <input_names> to None.')
inputs, input_names = list(inputs.values()), list(inputs.keys())
else:
inputs = nest.flatten(inputs)
# Process the outputs.
if isinstance(outputs, dict):
if output_names is not None:
raise ValueError(
'Excepted the output names from <outputs>.\n'
'You should set the <output_names> to None.')
outputs, output_names = list(outputs.values()), list(outputs.keys())
else:
outputs = nest.flatten(outputs)
if eager_context.executing_eagerly():
op_defs = []
tape = backprop.get_default_tape()
if tape is None:
raise RuntimeError('Please enter with ``onnx.frontend.record()``.')
for op_def in tape._defs:
op_defs.append(dragon_pb2.OperatorDef())
op_defs[-1].ParseFromString(op_def.SerializeAs())
graph_def = dragon_pb2.GraphDef(op=op_defs)
else:
symbolic_outputs = []
for output in outputs:
if types.is_symbolic_tensor(output):
symbolic_outputs.append(output)
graph_func = function_lib.create_function(outputs=symbolic_outputs)
graph_func.callback()
graph_def = graph_func.graph_def
graph_def.name = ''
# Add inputs and outputs.
for i, input in enumerate(inputs):
if hasattr(input, 'id'):
graph_def.input.extend([input.id])
elif input_names is not None:
graph_def.input.extend([input_names[i]])
for i, output in enumerate(outputs):
if hasattr(output, 'id'):
graph_def.output.extend([output.id])
elif output_names is not None:
graph_def.output.extend([output_names[i]])
# Make value info from inputs and outputs.
value_names = graph_def.input[:] + graph_def.output[:]
value_info = dict([(k, (helper.tensor_type(v.dtype), v.shape))
for k, v in zip(value_names, inputs + outputs)])
# Extract the constants from inputs and outputs.
constants = collections.OrderedDict()
for k, v in zip(value_names, inputs + outputs):
if isinstance(v, numpy.ndarray):
constants[k] = v
# Export.
model = graph_def_to_onnx_model(
graph_def=graph_def,
input_names=input_names,
output_names=output_names,
input_shapes=input_shapes,
constants=constants,
value_info=value_info,
opset_version=opset_version,
workspace=workspace_util.get_workspace(),
verbose=verbose,
enable_onnx_checker=enable_onnx_checker,
)
serialization.save_bytes(serialization.serialize_proto(model), f)
graph_def_to_onnx_graph = DragonFrontend.graph_def_to_onnx_graph
graph_def_to_onnx_model = DragonFrontend.graph_def_to_onnx_model
torch/core/onnx/utils.py → dragon/python/vm/onnx/core/frontend/torch.py
......@@ -8,6 +8,7 @@
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
"""PyTorch ONNX frontend."""
from __future__ import absolute_import
from __future__ import division
......@@ -19,7 +20,9 @@ import numpy
from dragon.core.framework import workspace
from dragon.core.proto import dragon_pb2
from dragon.core.util import nest
from dragon.vm.onnx.core import utils as onnx_util
from dragon.core.util import serialization
from dragon.vm.onnx.core import helper
from dragon.vm.onnx.core.frontend.native import graph_def_to_onnx_model
from dragon.vm.torch.core.autograd import backprop
......@@ -34,7 +37,7 @@ def export(
verbose=False,
enable_onnx_checker=True,
):
"""Export a model into ONNX format.
"""Export the recorded graph to an onnx model.
The outputs will be obtained by calling ``model(*args)``,
both the tensor or numpy array are allowed:
......@@ -162,10 +165,8 @@ def export(
# Make value info from inputs and outputs.
value_names = graph_def.input[:] + graph_def.output[:]
value_info = dict([
(k, onnx_util.make_value_info(v.shape, v.dtype))
for k, v in zip(value_names, inputs + outputs)
])
value_info = dict([(k, (helper.tensor_type(v.dtype), v.shape))
for k, v in zip(value_names, inputs + outputs)])
# Extract the constants from inputs and outputs.
constants = collections.OrderedDict()
......@@ -175,9 +176,8 @@ def export(
# Export.
with temporal_ws.as_default():
onnx_util.export_from_graph(
model = graph_def_to_onnx_model(
graph_def=graph_def,
f=f,
input_names=input_names,
output_names=output_names,
input_shapes=input_shapes,
......@@ -188,3 +188,4 @@ def export(
verbose=verbose,
enable_onnx_checker=enable_onnx_checker,
)
serialization.save_bytes(serialization.serialize_proto(model), f)
dragon/python/vm/onnx/core/shell.py deleted 100644 → 0
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
"""Toolkit for manipulating the onnx api."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import numpy
from dragon.core.autograph import function_lib
from dragon.core.eager import context
from dragon.core.eager import backprop
from dragon.core.framework import types
from dragon.core.framework import workspace
from dragon.core.proto import dragon_pb2
from dragon.core.util import nest
from dragon.vm.onnx.core import utils as onnx_util
class Shell(object):
"""Context-manger to export or load onnx models.
Enter a shell to export operators into an onnx model:
```python
x = dragon.constant([1, 2, 3])
with onnx.Shell() as shell, shell.as_default():
y = x * x
shell.export(inputs=[x], outputs=[y], f='model.onnx')
```
The onnx models can also be loaded to execute:
```python
f = shell.load_model('model.onnx', explicit_inputs=True)
print(f(np.array([1, 2, 3]))
```
"""
def __init__(self):
"""Create a ``Shell``."""
self._workspace = None
self._tape = backprop.GradientTape()
def as_default(self):
"""Set as the default shell."""
return self._workspace.as_default()
def export(
self,
inputs,
outputs,
f,
input_names=None,
output_names=None,
input_shapes=None,
opset_version=None,
verbose=False,
enable_onnx_checker=True,
):
"""Export an onnx model.
Parameters
----------
inputs : Union[Sequence, Dict]
The model inputs.
outputs : Union[Sequence, Dict]
The model outputs.
f : str
The filename for exporting model.
input_names : Sequence[str], optional
The name to the inputs.
output_names : Sequence[str], optional
The name to the outputs.
input_shapes : Union[Sequence, Dict], optional
The optional rewritten for input shapes.
opset_version : int, optional
The version of operator set.
verbose : bool, optional, default=False
Whether to print the debug string of graph.
enable_onnx_checker : bool, optional, default=True
Whether to check if model is valid.
"""
# Process the inputs.
if isinstance(inputs, dict):
if input_names is not None:
raise ValueError(
'Excepted the input names from <inputs>.\n'
'You should set the <input_names> to None.')
inputs, input_names = list(inputs.values()), list(inputs.keys())
else:
inputs = nest.flatten(inputs)
# Process the outputs.
if isinstance(outputs, dict):
if output_names is not None:
raise ValueError(
'Excepted the output names from <outputs>.\n'
'You should set the <output_names> to None.')
outputs, output_names = list(outputs.values()), list(outputs.keys())
else:
outputs = nest.flatten(outputs)
if context.executing_eagerly():
# Make graph def.
op_defs = []
for op_def in self._tape._tape._defs:
op_defs.append(dragon_pb2.OperatorDef())
op_defs[-1].ParseFromString(op_def.SerializeAs())
graph_def = dragon_pb2.GraphDef(op=op_defs)
else:
symbolic_outputs = []
for output in outputs:
if types.is_symbolic_tensor(output):
symbolic_outputs.append(output)
with self.as_default():
graph_func = function_lib.create_function(
outputs=symbolic_outputs)
graph_func.callback()
graph_def = graph_func.graph_def
graph_def.name = ''
# Add inputs and outputs.
for i, input in enumerate(inputs):
if hasattr(input, 'id'):
graph_def.input.extend([input.id])
elif input_names is not None:
graph_def.input.extend([input_names[i]])
for i, output in enumerate(outputs):
if hasattr(output, 'id'):
graph_def.output.extend([output.id])
elif output_names is not None:
graph_def.output.extend([output_names[i]])
# Make value info from inputs and outputs.
value_names = graph_def.input[:] + graph_def.output[:]
value_info = dict([
(k, onnx_util.make_value_info(v.shape, v.dtype))
for k, v in zip(value_names, inputs + outputs)
])
# Extract the constants from inputs and outputs.
constants = collections.OrderedDict()
for k, v in zip(value_names, inputs + outputs):
if isinstance(v, numpy.ndarray):
constants[k] = v
# Export.
onnx_util.export_from_graph(
graph_def=graph_def,
f=f,
input_names=input_names,
output_names=output_names,
input_shapes=input_shapes,
constants=constants,
value_info=value_info,
opset_version=opset_version,
workspace=self._workspace,
verbose=verbose,
enable_onnx_checker=enable_onnx_checker,
)
@staticmethod
def load_model(model_path, explicit_inputs=False):
"""Import an onnx model to the function.
Parameters
----------
model_path : str
The path to the onnx model.
explicit_inputs : bool, optional, default=False
**True** to attach model inputs to the function.
Returns
-------
callable
The function to run the model once.
"""
return onnx_util.import_to_function(model_path, explicit_inputs)
def __enter__(self):
self._workspace = workspace.Workspace()
self._workspace.merge_from(workspace.get_workspace())
if context.executing_eagerly():
self._tape._push_tape()
self._tape._tape.retain_graph = True
return self
def __exit__(self, typ, value, traceback):
if self._tape._recording:
self._tape._pop_tape()
dragon/python/vm/onnx/core/utils.py deleted 100644 → 0
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
"""Utilities for exporting and importing models."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy
from dragon.core.autograph import function_lib
from dragon.core.framework import workspace
from dragon.core.proto import dragon_pb2
from dragon.core.util import serialization
from dragon.vm.onnx.core.frontend import graph_def_to_onnx_model
from dragon.vm.onnx.core.helper import mapping
def export_from_graph(
graph_def,
f,
input_names=None,
output_names=None,
input_shapes=None,
constants=None,
value_info=None,
opset_version=None,
workspace=None,
verbose=True,
enable_onnx_checker=True,
):
"""Export an onnx model from the graph."""
model = graph_def_to_onnx_model(
graph_def=graph_def,
input_names=input_names,
output_names=output_names,
input_shapes=input_shapes,
constants=constants,
value_info=value_info,
opset_version=opset_version,
workspace=workspace,
verbose=verbose,
enable_onnx_checker=enable_onnx_checker)
serialization.save_bytes(serialization.serialize_proto(model), f)
def import_to_function(model_path, explicit_inputs=False):
"""Import an onnx model to the function."""
return function_lib \
.Function(name='onnx') \
.import_from(
graph_def=import_to_graph(model_path),
explicit_inputs=explicit_inputs,
)
def import_to_graph(model_path):
"""Import an onnx model to the graph."""
if not os.path.exists(model_path):
raise ValueError(
'Model({}) is not existed.'
.format(model_path)
)
graph_def = dragon_pb2.GraphDef()
serialized_proto = workspace \
.get_workspace() \
.ImportONNXModel(model_path)
graph_def.ParseFromString(serialized_proto)
return graph_def
def make_value_info(shape, dtype='float32'):
"""Return a value info from the shape and data type."""
return mapping.NP_TYPE_TO_TENSOR_TYPE[numpy.dtype(dtype)], shape
tensorflow/core/framework/constant_op.py
......@@ -50,6 +50,7 @@ def constant(value, dtype=None, shape=None, name='Const'):
The output tensor.
"""
dtype = str(dtype) if dtype else None
if dtype is not None:
if isinstance(value, numpy.ndarray):
value = value.astype(dtype)
......
tensorflow/core/framework/ops.py
......@@ -22,24 +22,17 @@ from dragon.core.framework import types
from dragon.vm.tensorflow.core.framework import constant_op
def convert_to_tensor(
value,
dtype=None,
name=None,
preferred_dtype=None,
):
def convert_to_tensor(value, dtype=None, name=None):
"""Converts the given value to a Tensor.
Parameters
----------
value : number, sequence or numpy.ndarray
value : Union[number, Sequence, numpy.ndarray]
The value to convert.
dtype : dragon.vm.tensorflow.dtypes.DType, optional
The optional data type.
name : str, optional
The Optional name.
preferred_dtype : dragon.vm.tensorflow.dtypes.DType, optional
The optional type when ``dtype`` is *None*.
Returns
-------
......@@ -108,4 +101,4 @@ def device(device_name):
id = int(id)
except Exception:
raise ValueError('The device id should be a integer.')
return context.device(device, device_id=id)
return context.device(device, device_index=id)
tensorflow/core/keras/engine/base_layer.py
......@@ -348,9 +348,7 @@ class Layer(module.Module):
else:
raise ValueError(
'Unknown format "%s".\n'
'Excepted format in (tf, h5, pkl).'
% (save_format,)
)
'Excepted format in (tf, h5, pkl).' % (save_format,))
if save_format == 'tf':
raise ValueError('TensorFlow format will never be supported.')
if save_format == 'h5':
......@@ -379,12 +377,9 @@ class Layer(module.Module):
'metrics',
}
if hasattr(self, '_layers'):
layers = layer_utils \
.filter_empty_layer_containers(self._layers)
return list(
itertools.chain.from_iterable(
getattr(layer, attribute) for layer in layers)
)
layers = layer_utils.filter_empty_layer_containers(self._layers)
return list(itertools.chain.from_iterable(
getattr(layer, attribute) for layer in layers))
return []
def _maybe_build(self, inputs):
......
tensorflow/core/keras/layers/convolutional.py
......@@ -163,7 +163,7 @@ class Conv2D(Conv):
The shape of convolution kernel.
strides : Sequence[int], optional, default=1
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding sizes.
data_format : {'channels_first', 'channels_last'}, optional
The optional data format.
......@@ -231,7 +231,7 @@ class Conv2DTranspose(Conv2D):
The shape of convolution kernel.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding sizes.
output_padding : Sequence[int], optional
The sizes of padded to the output.
......@@ -396,7 +396,7 @@ class DepthwiseConv2D(Conv2D):
The shape of convolution kernel.
strides : Sequence[int], optional, default=1
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding sizes.
data_format : {'channels_first', 'channels_last'}, optional
The optional data format.
......
tensorflow/core/keras/layers/pooling.py
......@@ -106,7 +106,7 @@ class AveragePooling2D(Pooling2D):
The size(s) of sliding window.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding sizes.
data_format : {'channels_first', 'channels_last'}, optional
The optional data format.
......@@ -179,7 +179,7 @@ class MaxPooling2D(Pooling2D):
The size(s) of sliding window.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding sizes.
data_format : {'channels_first', 'channels_last'}, optional
The optional data format.
......
tensorflow/core/ops/nn_ops.py
......@@ -44,7 +44,7 @@ def avg_pool(
The size(s) of sliding window.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......@@ -104,7 +104,7 @@ def avg_pool2d(
The size(s) of sliding window.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......@@ -157,7 +157,7 @@ def convolution(
The weight tensor.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding size(s).
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......@@ -223,7 +223,7 @@ def conv_transpose(
The determined shape of output.
strides : Sequence[int]
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]], optional
padding : Union[str, Sequence[int]], optional
The padding algorithm or padding size(s).
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......@@ -292,7 +292,7 @@ def conv2d(
The weight tensor.
strides : Sequence[int]
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NHWC'}, optional
The optional data format.
......@@ -332,7 +332,7 @@ def conv2d_transpose(
The determined shape of output.
strides : Sequence[int]
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NHWC'}, optional
The optional data format.
......@@ -370,7 +370,7 @@ def depthwise_conv2d(
The weight tensor.
strides : Sequence[int]
The stride(s) of sliding window.
padding : Union[{'VALID', 'SAME'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NHWC'}, optional
The optional data format.
......@@ -590,7 +590,7 @@ def max_pool(
The size(s) of sliding window.
strides : Sequence[int]
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding sizes.
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......@@ -650,7 +650,7 @@ def max_pool2d(
The size(s) of sliding window.
strides : Sequence[int], optional
The stride(s) of sliding window.
padding : Union[{'valid', 'same'}, Sequence[int]]
padding : Union[str, Sequence[int]]
The padding algorithm or padding size(s).
data_format : {'NCHW', 'NCDHW', 'NHWC', 'NDHWC'}, optional
The optional data format.
......
tensorflow/core/ops/variables.py
......@@ -33,6 +33,8 @@ class VariableMetaclass(object):
class Variable(VariableMetaclass, EagerTensor):
"""Resource variable."""
def __init__(
self,
initial_value,
......@@ -44,6 +46,7 @@ class Variable(VariableMetaclass, EagerTensor):
"""Create a ``Variable``."""
super(Variable, self).__init__(trainable=trainable)
name = name if name else 'Variable'
dtype = str(dtype) if dtype else None
self._name = context.get_name_scope() + name + ':0'
# Determine th value.
if isinstance(initial_value, EagerTensor):
......@@ -52,7 +55,7 @@ class Variable(VariableMetaclass, EagerTensor):
initial_value = initial_value.get_value()
# Determine the data type.
if not isinstance(initial_value, numpy.ndarray):
initial_value = numpy.array(initial_value, dtype if dtype else dtype)
initial_value = numpy.array(initial_value, dtype)
elif dtype is not None:
initial_value = initial_value.astype(dtype)
# Determine the tensor shape.
......@@ -103,8 +106,7 @@ def get_default_initializer(name, shape=None, dtype=dtypes.float32):
else:
raise ValueError(
'An initializer for Variable({}) of {} is required.'
.format(name, dtype.base_dtype)
)
.format(name, dtype.base_dtype))
return initializer
......
tensorlayer/core/engine/model.py
......@@ -3,6 +3,7 @@ from __future__ import division
from __future__ import print_function
from dragon.core.framework import context
from dragon.core.util import string
from dragon.vm.tensorlayer.core.engine import module
......@@ -83,10 +84,10 @@ class Model(module.Module):
return self.forward(inputs, **kwargs)
def __repr__(self):
tmpstr = self.name + '(\n'
tmp_str = self.name + '(\n'
for idx, layer in enumerate(self.all_layers):
modstr = layer.__repr__()
modstr = self._addindent(modstr, 2)
tmpstr = tmpstr + ' (' + layer.name + '): ' + modstr + '\n'
tmpstr = tmpstr + ')'
return tmpstr
mod_str = layer.__repr__()
mod_str = string.add_indent(mod_str, 2)
tmp_str = tmp_str + ' (' + layer.name + '): ' + mod_str + '\n'
tmp_str = tmp_str + ')'
return tmp_str
tensorlayer/core/engine/module.py
......@@ -201,15 +201,13 @@ class Module(object):
else:
raise ValueError(
"Saving format should be in ('hdf5', 'npz', 'pkl', 'npz_dict').\n"
"Format <%s> is not supported." % format
)
"Format <%s> is not supported." % format)
if verbose:
for info in matched_info:
logging.info(
'Weight({}) loaded, Size: ({})'
.format(info[0], ', '.join([str(d) for d in info[1]]))
)
.format(info[0], ', '.join([str(d) for d in info[1]])))
def save_weights(self, filepath, format=None):
"""Save weights into a binary file.
......@@ -241,8 +239,7 @@ class Module(object):
else:
raise ValueError(
"Saving format should be in ('hdf5', 'npz', 'pkl', 'npz_dict').\n"
"Format <%s> is not supported." % format
)
"Format <%s> is not supported." % format)
@staticmethod
def _dedupe_weights(weights):
......
tensorrt/python/__init__.py
......@@ -18,15 +18,14 @@ import os as _os
import sys as _sys
# Modules
from dragon.vm.tensorrt._api import backend
from dragon.vm.tensorrt._api import onnx
# Classes
from dragon.vm.tensorrt.core.backend import ONNXBackendRep
from dragon.vm.tensorrt.core.engine import Binding
from dragon.vm.tensorrt.core.engine import Engine
# Attributes
_API_MODULE = backend
_API_MODULE = onnx
_current_module = _sys.modules[__name__]
_api_dir = _os.path.dirname(_os.path.dirname(_API_MODULE.__file__))
if not hasattr(_current_module, '__path__'):
......
tensorrt/python/_api/onnx/__init__.py 0 → 100644
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
from __future__ import absolute_import as _absolute_import
from __future__ import division as _division
from __future__ import print_function as _print_function
# Classes
from dragon.vm.onnx.core.backend.tensorrt import BackendRep
# Functions
from dragon.vm.onnx.core.backend.tensorrt import prepare as prepare_backend
from dragon.vm.onnx.core.backend.tensorrt import run_model
from dragon.vm.onnx.core.backend.tensorrt import run_node
from dragon.vm.onnx.core.backend.tensorrt import supports_device
# Attributes
__all__ = [_s for _s in dir() if not _s.startswith('_')]
tensorrt/python/core/engine.py
......@@ -442,12 +442,10 @@ def check_input_validity(index, array, binding):
raise TypeError(
'Wrong dtype for input %i.\n'
'Expected %s, got %s. Cannot safely cast.' %
(index, binding.dtype, array.dtype)
)
(index, binding.dtype, array.dtype))
else:
raise TypeError(
'Wrong dtype for input %i.\n'
'Expected %s, got %s.' %
(index, binding.dtype, array.dtype)
)
(index, binding.dtype, array.dtype))
return array
torch/_api/onnx/__init__.py
......@@ -14,6 +14,6 @@ from __future__ import absolute_import as _absolute_import
from __future__ import division as _division
from __future__ import print_function as _print_function
from dragon.vm.torch.core.onnx.utils import export
from dragon.vm.onnx.core.frontend.torch import export
__all__ = [_s for _s in dir() if not _s.startswith('_')]
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!