| """Linux CPU exporter; use an isolated venv, never the training environment.""" |
| import os |
| os.environ.update(CUDA_VISIBLE_DEVICES='',OMP_NUM_THREADS='2',OPENBLAS_NUM_THREADS='1',TF_NUM_INTRAOP_THREADS='2',TF_NUM_INTEROP_THREADS='1') |
| import argparse,json,time,sys |
| from pathlib import Path |
| import torch |
| import litert_torch |
| from ai_edge_quantizer import quantizer,recipe |
| from build_recognizer import load,R |
| from portable_model import PortableRecognizer,Compact |
|
|
| def main(): |
| parser=argparse.ArgumentParser();parser.add_argument('--detector',action='store_true');parser.add_argument('--checkpoint',type=Path);args=parser.parse_args() |
| torch.set_num_threads(2);out=R/'litert';out.mkdir(exist_ok=True) |
| if args.detector: |
| sys.path.append('/home/rezo/bina02-surya-layout-env/lib/python3.10/site-packages') |
| from export_detector import EfficientViTConfig,EfficientViTForSemanticSegmentation,BASE,S,Model |
| m=EfficientViTForSemanticSegmentation(EfficientViTConfig.from_pretrained(str(BASE))).eval() |
| m.load_state_dict(torch.load(args.checkpoint or S/'run/checkpoints/step-001000.pt',map_location='cpu',weights_only=True)['model']) |
| converter=litert_torch.signature('detector',Model(m).eval(),(torch.zeros(1,3,1200,1200),)) |
| prefix='detector' |
| else: |
| model=Compact(PortableRecognizer(load(args.checkpoint or R/'source/recognizer-095000.pt'))).eval() |
| model.model.index_dtype=torch.int32 |
| converter=None |
| for width in [512,1024,2048,4096,8192]: |
| x=torch.zeros(1,1,64,width);w=torch.tensor([width]) |
| if converter is None:converter=litert_torch.signature('width'+str(width),model,(x,w)) |
| else:converter.signature('width'+str(width),model,(x,w)) |
| prefix='recognizer' |
| with torch.no_grad():edge=converter.convert() |
| fp=out/(prefix+'-fp32.tflite');edge.export(str(fp)) |
| results=[] |
| for name,rp in [('weight-only4',recipe.weight_only_wi4_afp32()),('dynamic4-block32',recipe.dynamic_wi4b32_afp32())]: |
| t=time.monotonic() |
| try: |
| q=quantizer.Quantizer(str(fp));q.load_quantization_recipe(rp) |
| target=out/f'{prefix}-{name}.tflite';q.quantize().export_model(str(target)) |
| results.append(dict(name=name,path=str(target),bytes=target.stat().st_size,recipe=rp,seconds=time.monotonic()-t)) |
| except Exception as e:results.append(dict(name=name,error=repr(e))) |
| (out/(prefix+'-build.json')).write_text(json.dumps(dict(converter='litert-torch 0.10.0',quantizer='ai-edge-quantizer 0.10.0', |
| input='Five static width signatures512/1024/2048/4096/8192; true widths input; zero right pad, never truncate',results=results),indent=2,default=str)+'\n') |
| print(json.dumps(results,default=str),flush=True) |
|
|
| if __name__=='__main__':main() |
|
|
