1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
|
from torch.onnx.symbolic_helper import parse_args
import torch.onnx.symbolic_helper as sym_help
import torch.onnx.symbolic_registry as sym_registry
import importlib
from inspect import getmembers, isfunction
def register_quantized_ops(domain, version):
# Register all the non-quantized ops
sym_registry.register_version('', version)
# Register all quantized ops
module = importlib.import_module('torch.onnx.symbolic_caffe2')
sym_registry._symbolic_versions['caffe2'] = module
quant_version_ops = getmembers(sym_registry._symbolic_versions['caffe2'])
for op in quant_version_ops:
if isfunction(op[1]) and not sym_registry.is_registered_op(op[0], domain, version):
aten_q_ops = ['relu', '_empty_affine_quantized', 'dequantize',
'quantize_per_tensor', 'upsample_nearest2d', 'avg_pool2d',
'reshape', 'slice', 'cat', 'max_pool2d', 'sigmoid']
if op[0] in aten_q_ops:
sym_registry.register_op(op[0], op[1], '', version)
sym_registry.register_op(op[0], op[1], domain, version)
def _permute_helper(g, input, axes):
quant_args = {
"axes_i": axes,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Transpose", input, **quant_args)
sym_help._quantized_ops.add(output)
return output
def nchw2nhwc(g, input):
axes = [0, 2, 3, 1]
return _permute_helper(g, input, axes)
def nhwc2nchw(g, input):
axes = [0, 3, 1, 2]
return _permute_helper(g, input, axes)
def linear_prepack(g, weight, bias):
# Mapping to a dummy caffe2 prepack node.
# During the onnx -> c2 conversion we can look up original weight and bias
# from this node
output = g.op("_caffe2::WeightPrepack", weight, bias)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'v', 'v', 'f', 'i')
def linear(g, input, weight, bias, scale, zero_point):
kwargs = {
"Y_scale_f": scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8FC", input, weight, bias, **kwargs)
sym_help._quantized_ops.add(output)
return output
def conv_prepack(g, input, weight, bias, stride, padding, dilation, groups):
# Mapping to a dummy caffe2 prepack node.
# During the onnx -> c2 conversion we can look up original weight and bias
# from this node
output = g.op("_caffe2::WeightPrepack", input, weight, bias)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'v', 'v', 'is', 'is', 'is', 'i', 'f', 'i')
def conv2d(g, input, weight, bias, stride, padding, dilation, groups, scale, zero_point):
kernel_size = weight.node()["shape"][1:3]
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"dilations_i": dilation,
"group_i": groups,
"kernels_i": kernel_size,
"order_s": "NHWC",
"Y_scale_f": scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8Conv", input, weight, bias, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'v', 'v', 'is', 'is', 'is', 'i', 'f', 'i')
def conv2d_relu(g, input, weight, bias, stride, padding, dilation, groups, scale, zero_point):
kernel_size = weight.node()["shape"][1:3]
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"dilations_i": dilation,
"group_i": groups,
"kernels_i": kernel_size,
"order_s": "NHWC",
"Y_scale_f": scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8ConvRelu", input, weight, bias, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'v', 'f', 'i')
def add(g, input_a, input_b, scale, zero_point):
kwargs = {
"Y_scale_f": scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8Add", input_a, input_b, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v')
def relu(g, input):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import relu
return relu(g, input)
kwargs = {
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Relu", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'f', 'i', 't')
def quantize_per_tensor(g, input, scale, zero_point, dtype):
kwargs = {
"Y_scale_f": scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8Quantize", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v')
def dequantize(g, input):
return g.op("_caffe2::Int8Dequantize", input)
@parse_args('v', 't', 't', 't', 't', 't', 't', 't')
def _empty_affine_quantized(g, input, shape, scale, zero_point, dtype, pin_memory, memory_format, layout):
return input
def upsample_nearest2d(g, input, output_size, align_corners=None, scales_h=None, scales_w=None):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import upsample_nearest2d as upsample_nearest2d_impl
return upsample_nearest2d_impl(g, input, output_size, align_corners)
output_size = sym_help._parse_arg(output_size, 'is')
kwargs = {
"output_size_i": output_size,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8ResizeNearest", input, **kwargs)
output = nhwc2nchw(g, output)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'is', 'is', 'is', 'is', 'i')
def max_pool2d(g, input, kernel_size, stride, padding, dilation, ceil_mode):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import max_pool2d
return max_pool2d(g, input, kernel_size, stride, padding, dilation, ceil_mode)
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"kernel_i": kernel_size[0],
"order_s": "NHWC",
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8MaxPool", input, **kwargs)
output = nhwc2nchw(g, output)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'is', 'is', 'is', 'i', 'i', 'none')
def avg_pool2d(g, input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override=None):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import avg_pool2d
return avg_pool2d(g, input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)
kwargs = {
"strides_i": stride,
"pads_i": padding + padding,
"kernel_i": kernel_size[0],
"order_s": "NHWC",
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8AveragePool", input, **kwargs)
output = nhwc2nchw(g, output)
sym_help._quantized_ops.add(output)
return output
def reshape(g, input, shape):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import reshape
return reshape(g, input, shape)
kwargs = {
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Reshape", input, shape, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v', 'v', 'v', 'v', 'i')
def slice(g, input, dim, start, end, step):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import slice
return slice(g, input, dim, start, end, step)
if step != 1:
raise RuntimeError("ONNX quantized slice export only works for step 1.")
start = sym_help._parse_arg(start, 'i')
end = sym_help._parse_arg(end, 'i')
dim = sym_help._parse_arg(dim, 'i')
kwargs = {
"start_idx_i": start,
"end_idx_i": end,
"dim_i": dim,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Slice", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
def cat(g, tensor_list, dim, scale=None, zero_point=None):
tensors = sym_help._unpack_list(tensor_list)
input = tensors[0]
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import cat
return cat(g, tensor_list, dim)
dim = sym_help._parse_arg(dim, 'i')
kwargs = {
"Y_scale_f": tensors[0].node()["Y_scale"],
"Y_zero_point_i": tensors[0].node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Concat", *tensors, axis_i=dim, **kwargs)
sym_help._quantized_ops.add(output)
return output
@parse_args('v')
def sigmoid(g, input):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import sigmoid
return sigmoid(g, input)
# Caffe2 expects the output scale to be 1/2^8
# and output zero_point to be 0 (quint8 type)
out_scale = 1.0 / 256
zero_point = 0
kwargs = {
"Y_scale_f": out_scale,
"Y_zero_point_i": zero_point,
}
output = g.op("_caffe2::Int8Sigmoid", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
|
