File size: 1,783 Bytes
6986878 | 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 | """Portable INT8 weight storage with FP32 arithmetic (no dynamic activation quantization)."""
import numpy as np
import onnx
from onnx import numpy_helper, helper, TensorProto
def quantize_weights(source, destination):
model = onnx.load(str(source))
initializers = []
dequantizers = []
gather_weights = {node.input[0] for node in model.graph.node if node.op_type == 'Gather'}
for tensor in model.graph.initializer:
if tensor.data_type != TensorProto.FLOAT or len(tensor.dims) != 2:
initializers.append(tensor)
continue
array = numpy_helper.to_array(tensor)
axis = 0 if tensor.name in gather_weights else 1
scale = np.maximum(np.max(np.abs(array), axis=1-axis) / 127, np.finfo(np.float32).tiny)
expanded = scale[:, None] if axis == 0 else scale[None, :]
quantized = np.clip(np.rint(array / expanded), -127, 127).astype(np.int8)
weight, factor, zero = [tensor.name + suffix for suffix in ('_int8', '_scale', '_zero')]
initializers.extend([numpy_helper.from_array(quantized, weight),
numpy_helper.from_array(scale, factor),
numpy_helper.from_array(np.zeros_like(scale, dtype=np.int8), zero)])
dequantizers.append(helper.make_node('DequantizeLinear', [weight, factor, zero],
[tensor.name], axis=axis, name=tensor.name+'_dequantize'))
del model.graph.initializer[:]
model.graph.initializer.extend(initializers)
nodes = list(model.graph.node)
del model.graph.node[:]
model.graph.node.extend(dequantizers + nodes)
onnx.save_model(model, str(destination))
if __name__ == '__main__':
import sys
quantize_weights(sys.argv[1], sys.argv[2])
|