| |
| |
| |
| |
| |
| import sys |
| import numpy as np |
| import onnx |
| from onnx import TensorProto, helper, numpy_helper |
| from onnxruntime.quantization.matmul_nbits_quantizer import DefaultWeightOnlyQuantConfig, MatMulNBitsQuantizer |
|
|
| src, dst = sys.argv[1], sys.argv[2] |
| m = onnx.load(src) |
| q = MatMulNBitsQuantizer(m, algo_config=DefaultWeightOnlyQuantConfig(block_size=32, is_symmetric=True, bits=8, op_types_to_quantize=('MatMul',), quant_axes=(('MatMul', 0),))) |
| q.process() |
| m = q.model.model |
| g = m.graph |
| inits = {i.name: i for i in g.initializer} |
| node = next(n for n in g.node if n.op_type == 'Gather' and n.input[0] in inits and len(inits[n.input[0]].dims) == 2 and inits[n.input[0]].dims[0] > 100000) |
| table = numpy_helper.to_array(inits[node.input[0]]).astype(np.float32) |
| scale = np.maximum(np.abs(table).max(axis=1, keepdims=True), 1e-12) / 127.0 |
| rows = np.clip(np.rint(table / scale), -127, 127).astype(np.int8) |
| name, ids, out = node.input[0], node.input[1], node.output[0] |
| g.initializer.remove(inits[name]) |
| g.initializer.extend([numpy_helper.from_array(rows, name + '_int8'), numpy_helper.from_array(scale.astype(np.float32), name + '_scale')]) |
| at = list(g.node).index(node) |
| g.node.remove(node) |
| for i, n in enumerate([ |
| helper.make_node('Gather', [name + '_int8', ids], [out + '_q'], axis=0, name='table_gather_int8'), |
| helper.make_node('Cast', [out + '_q'], [out + '_f'], to=TensorProto.FLOAT, name='table_cast'), |
| helper.make_node('Gather', [name + '_scale', ids], [out + '_s'], axis=0, name='table_gather_scale'), |
| helper.make_node('Mul', [out + '_f', out + '_s'], [out], name='table_dequantize'), |
| ]): |
| g.node.insert(at + i, n) |
| onnx.checker.check_model(m) |
| onnx.save(m, dst) |
| print('saved', dst) |
|
|
