File size: 5,638 Bytes
618450b | 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 | 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())
# Near-empty synthetic crops amplify fp32 summation-order noise in SE.
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
# Dynamic width with exact mask semantics; small bounded batch for mobile.
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()
|