| import argparse, hashlib, json, sys, time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from portable_model import PortableRecognizer, Compact |
|
|
| R = Path(__file__).resolve().parent |
| sys.path.insert(0, str(R/'source')) |
| from recognizer import Recognizer |
|
|
| def load(path): |
| state = torch.load(path, map_location='cpu', weights_only=True)['model'] |
| source = Recognizer(classes=state['head.weight'].shape[0], gradient_checkpointing=False).eval() |
| source.load_state_dict(state) |
| return source |
|
|
| def main(): |
| p=argparse.ArgumentParser() |
| p.add_argument('--checkpoint',type=Path,default=R/'source/recognizer-095000.pt') |
| p.add_argument('--out',type=Path,default=R/'models') |
| p.add_argument('--coreml',action='store_true') |
| p.add_argument('--onnx',action='store_true') |
| p.add_argument('--hqq',action='store_true') |
| p.add_argument('--blocks',default='16,32,64') |
| a=p.parse_args();a.out.mkdir(parents=True,exist_ok=True) |
| torch.set_num_threads(2);torch.manual_seed(18404) |
| source=load(a.checkpoint);model=PortableRecognizer(source).eval() |
| controls=[] |
| with torch.inference_mode(): |
| for widths in [[1,7,64],[63,257,511],[769,1025,1537]]: |
| w=torch.tensor(widths,dtype=torch.int64);x=torch.rand(len(w),1,64,((max(widths)+7)//8)*8) |
| for i,n in enumerate(widths):x[i,:,:,n:]=0 |
| y,n=source(x,w);z,m=model(x,w) |
| error=float((y-z).abs().max()); agreement=float((y.argmax(-1)==z.argmax(-1)).float().mean()) |
| |
| tolerance = .02 if min(widths)<8 else 2e-4 |
| assert error < tolerance and agreement == 1.0 and torch.equal(n,m), (widths,error,agreement) |
| controls.append(dict(widths=widths,max_abs_logit_error=error,frame_argmax_agreement=agreement)) |
| identity=dict(checkpoint=str(a.checkpoint.resolve()),sha256=hashlib.sha256(a.checkpoint.read_bytes()).hexdigest(), |
| parameters=sum(t.numel()for t in source.parameters()),lowering_controls=controls, |
| torch=torch.__version__,output_contract=['tokens:int32','confidence:float32','lengths:int32']) |
| (a.out/'source.json').write_text(json.dumps(identity,indent=2)+'\n') |
| compact=Compact(model).eval();x=torch.zeros(1,1,64,512);w=torch.tensor([512]) |
| if a.onnx: |
| import onnx |
| from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer,DefaultWeightOnlyQuantConfig |
| from onnxruntime.quantization.quant_utils import load_model_with_shape_infer |
| fp=a.out/'recognizer-fp32.onnx' |
| torch.onnx.export(compact,(x,w),str(fp),opset_version=18,dynamo=False, |
| input_names=['pixels','widths'],output_names=['tokens','confidence','lengths'], |
| dynamic_axes={'pixels':{0:'batch',3:'width'},'widths':{0:'batch'}, |
| 'tokens':{0:'batch',1:'frames'},'confidence':{0:'batch',1:'frames'},'lengths':{0:'batch'}}) |
| graph=onnx.load(fp);onnx.checker.check_model(graph) |
| for block in map(int,a.blocks.split(',')): |
| q=MatMulNBitsQuantizer(load_model_with_shape_infer(fp), |
| algo_config=DefaultWeightOnlyQuantConfig(block_size=block,is_symmetric=True,accuracy_level=4)) |
| q.process();target=a.out/f'recognizer-w4-block{block}.onnx' |
| q.model.save_model_to_file(str(target),use_external_data_format=False) |
| graph=onnx.load(target);onnx.checker.check_model(graph) |
| counts={} |
| for node in graph.graph.node:counts[node.op_type]=counts.get(node.op_type,0)+1 |
| (target.with_suffix('.audit.json')).write_text(json.dumps(dict(bytes=target.stat().st_size, |
| ops=counts,weight_bits=4,graph_input_bits=32,accuracy_level=4,retained_float_ops=['stem','depthwise convolutions','biases'], |
| format='ORT MatMulNBits W4 with runtime-dependent internal precision; mixed-weight precision, not all-INT4'),indent=2)+'\n') |
| if a.hqq: |
| from onnxruntime.quantization.matmul_nbits_quantizer import HQQWeightOnlyQuantConfig |
| q=MatMulNBitsQuantizer(load_model_with_shape_infer(fp),algo_config=HQQWeightOnlyQuantConfig(block_size=32,bits=4)) |
| q.process();q.model.save_model_to_file(str(a.out/'recognizer-hqq4-block32.onnx'),use_external_data_format=False) |
| if a.coreml: |
| import coremltools as ct |
| from coremltools.optimize.coreml import OptimizationConfig,OpLinearQuantizerConfig,linear_quantize_weights |
| |
| traced=torch.jit.trace(compact,(x,w)) |
| ml=ct.convert(traced,inputs=[ct.TensorType(name='pixels',shape=(ct.RangeDim(1,8,default=1),1,64,ct.RangeDim(8,8192,default=512))), |
| ct.TensorType(name='widths',shape=(ct.RangeDim(1,8,default=1),),dtype=np.int32)], |
| outputs=[ct.TensorType(name='tokens'),ct.TensorType(name='confidence'),ct.TensorType(name='lengths')], |
| minimum_deployment_target=ct.target.iOS18,compute_precision=ct.precision.FLOAT32, |
| compute_units=ct.ComputeUnit.CPU_AND_GPU) |
| ml.save(str(a.out/'recognizer-fp32.mlpackage')) |
| for block in map(int,a.blocks.split(',')): |
| config=OptimizationConfig(global_config=OpLinearQuantizerConfig(mode='linear_symmetric',dtype='int4', |
| granularity='per_block',block_size=block,weight_threshold=512)) |
| quant=linear_quantize_weights(ml,config=config) |
| quant.save(str(a.out/f'recognizer-w4-block{block}.mlpackage')) |
| print(json.dumps(identity),flush=True) |
|
|
| if __name__=='__main__':main() |
|
|
