| """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]) |
|
|
