Skip to content

Commit 23543f2

Browse files
Laicheng0830Laicheng0830
authored
TLX2ONNX 0.0.1 release
1 parent 91eabc3 commit 23543f2

File tree

6 files changed

+145
-19
lines changed

6 files changed

+145
-19
lines changed

‎OP_LIST.md‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,17 +53,27 @@
5353
|Tile| 1~12|Supported|
5454
|UpSampling2d| 1~12|Supported|
5555
|DownSampling2d| 1~12|Supported|
56-
|Concat|
57-
Elementwise|
58-
|GaussianNoise|
59-
|PadLayer|
60-
|ZeroPad1d|
61-
|ZeroPad2d|
62-
|ZeroPad3d|
63-
|Stack|
64-
|UnStack|
65-
|Sign|
66-
|Scale|
56+
|Concat| 1~12 | Supported|
57+
|Elementwise| 1~12 | Supported|
58+
|GaussianNoise| 1~12 | Supported|
59+
|PadLayer| 1~12 | Supported|
60+
|ZeroPad1d| 1~12 | Supported|
61+
|ZeroPad2d| 1~12 | Supported|
62+
|ZeroPad3d| 1~12 | Supported|
63+
|Stack| 1~12 | Supported|
64+
|UnStack| 1~12 | Supported|
65+
|Scale| 1~12 | Supported|
66+
|RNN|1~12 | Supported|
67+
|RNNCell|1~12 | Supported|
68+
|LSTM|1~12 | Supported|
69+
|LSTMCell|1~12 | Supported|
70+
|GRU|1~12 | Supported|
71+
|GRUCell|1~12 | Supported|
72+
|LayerNorm| 17 | Supported|
73+
|GroupConv2d
74+
|SeparableConv1d
75+
|SeparableConv2d
76+
|SubpixelConv2d
6777
| Matmul | 1~12 | Supported|
6878

6979

@@ -75,3 +85,5 @@ Elementwise|
7585

7686

7787

88+
89+

‎requirements.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
onnx<=1.11.0
2-
tensorlayerx>=0.5.1
2+
tensorlayerx>=0.5.6

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
MAJOR = 0
99
MINOR = 0
1010
PATCH = 1
11-
PRE_RELEASE = 'alpha'
11+
PRE_RELEASE = ''
1212
# Use the following formatting: (major, minor, patch, prerelease)
1313
VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)
1414

‎tests/test_layernorm.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#! /usr/bin/python
2+
# -*- coding: utf-8 -*-
3+
4+
import os
5+
os.environ["TL_BACKEND"] = 'tensorflow'
6+
import tensorlayerx as tlx
7+
from tensorlayerx.nn import Module
8+
from tensorlayerx.nn import LayerNorm
9+
from tlx2onnx.main import export
10+
import onnxruntime as rt
11+
import numpy as np
12+
13+
14+
class NET(Module):
15+
def __init__(self):
16+
super(NET, self).__init__()
17+
self.layernorm = LayerNorm([50, 50, 32])
18+
19+
def forward(self, x):
20+
x = self.layernorm(x)
21+
return x
22+
23+
net = NET()
24+
print(type(net))
25+
net.set_eval()
26+
input = tlx.nn.Input(shape=(10, 50, 50, 32))
27+
onnx_model = export(net, input_spec=input, path='layernorm.onnx', enable_onnx_checker=False)
28+
print("tlx out", input)
29+
30+
# Infer Model
31+
sess = rt.InferenceSession('layernorm.onnx')
32+
33+
input_name = sess.get_inputs()[0].name
34+
output_name = sess.get_outputs()[0].name
35+
36+
input_data = np.array(input, dtype=np.float32)
37+
38+
result = sess.run([output_name], {input_name: input_data})
39+
print("onnx out", result)

‎tlx2onnx/main.py‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,18 +9,33 @@
99
from .common import make_graph, logging
1010
from .op_mapper.datatype_mapping import NP_TYPE_TO_TENSOR_TYPE
1111

12-
def export(model, input_spec, path=None, export_params=False, opset_version = 9, auto_update_opset=True):
12+
def export(model, input_spec, path=None, enable_onnx_checker=True, opset_version = 9, auto_update_opset=True):
1313
"""
1414
1515
Parameters
1616
----------
17-
model
18-
input_spec
19-
path
20-
export_params
17+
model : object
18+
TensorLayerX instantiate the net object.
19+
input_spec : tensor
20+
TensorLayerX Input.
21+
path : string
22+
ONNX file saving path
23+
enable_onnx_checker : bool
24+
Whether to enable ONNX model checker.
25+
opset_version : int
26+
The version of the default (ai.onnx) opset to target. Must be >= 7 and <= 17.
2127
2228
Returns
2329
-------
30+
ONNX model file
31+
32+
Examples
33+
---------
34+
>>> class NET(Module):
35+
>>> net = NET()
36+
>>> net.set_eval()
37+
>>> input = tlx.nn.Input([10, 50, 50, 32], name='input')
38+
>>> onnx_model = export(net, input_spec=input, path='vgg.onnx')
2439
2540
"""
2641

@@ -64,7 +79,8 @@ def export(model, input_spec, path=None, export_params=False, opset_version = 9,
6479
producer_name='onnx-mode'
6580
)
6681

67-
onnx.checker.check_model(model_def)
82+
if enable_onnx_checker:
83+
onnx.checker.check_model(model_def)
6884
onnx.save(model_def, path)
6985
logging.info("ONNX model saved in {}".format(path))
7086
return model_def

‎tlx2onnx/op_mapper/nn/normalization.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,63 @@ def version_1(cls, node, **kwargs):
7979
outputs=node['out_nodes_name']
8080
)
8181
onnx_node.append(bn_node)
82+
return onnx_node, onnx_value, onnx_init
83+
84+
85+
@OpMapper(['LayerNorm'])
86+
class LayerNorm():
87+
# supports v17
88+
89+
@classmethod
90+
def version_17(cls, node, **kwargs):
91+
onnx_node = []
92+
onnx_value = []
93+
onnx_init = []
94+
95+
# input , output, data_format
96+
x = node['in_nodes_name'][0]
97+
x_shape = node['in_tensors'][0]
98+
99+
out_shape = node['out_tensors'][0]
100+
out_v = helper.make_tensor_value_info(node['out_nodes_name'][0], NP_TYPE_TO_TENSOR_TYPE[node['dtype']],
101+
shape=node['out_tensors'][0])
102+
onnx_value.append(out_v)
103+
104+
spatial = 2
105+
# get parameters
106+
beta_name = node['node'].layer.name + '/beta'
107+
beta_weight = numpy_helper.from_array(arr=to_numpy(node['node'].layer.beta), name=beta_name)
108+
onnx_init.append(beta_weight)
109+
110+
gamma_name = node['node'].layer.name + '/gamma'
111+
gamma_weight = numpy_helper.from_array(arr=to_numpy(node['node'].layer.gamma), name=gamma_name)
112+
onnx_init.append(gamma_weight)
113+
114+
epsilon = node['node'].layer.epsilon
115+
116+
# if data_format == 'channels_last':
117+
# channels last conver weights and input
118+
x_shape = make_shape_channels_first(x_shape)
119+
out_temp_shape = make_shape_channels_first(out_shape)
120+
# make channels transpose
121+
t_x = helper.make_tensor_value_info(node['in_nodes_name'][0] + 't',
122+
NP_TYPE_TO_TENSOR_TYPE[node['dtype']], shape=x_shape)
123+
onnx_value.append(t_x)
124+
tx_node, x = make_node('Transpose', inputs=[x], outputs=[node['in_nodes_name'][0] + 't'],
125+
perm=get_channels_first_permutation(spatial))
126+
onnx_node.append(tx_node)
127+
# make batch normalization
128+
out_temp = helper.make_tensor_value_info(node['out_nodes_name'][0] + 'bn',
129+
NP_TYPE_TO_TENSOR_TYPE[node['dtype']], shape=out_temp_shape)
130+
onnx_value.append(out_temp)
131+
ln_node, out = make_node('LayerNormalization',
132+
inputs=[node['in_nodes_name'][0] + 't', beta_name, gamma_name],
133+
outputs=[node['out_nodes_name'][0] + 'bn'], epsilon=epsilon
134+
)
135+
onnx_node.append(ln_node)
136+
# make channels transpose
137+
t_out = helper.make_tensor_value_info(node['out_nodes_name'][0], NP_TYPE_TO_TENSOR_TYPE[node['dtype']], shape=out_shape)
138+
onnx_value.append(t_out)
139+
tout_node, _ = make_node('Transpose', inputs=[out], outputs=node['out_nodes_name'], perm=get_channels_last_permutation(spatial))
140+
onnx_node.append(tout_node)
82141
return onnx_node, onnx_value, onnx_init

0 commit comments

Comments
 (0)