diff --git a/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json new file mode 100644 index 000000000..7713b8e7d --- /dev/null +++ b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json @@ -0,0 +1,61 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 256, + 192 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "heatmaps" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "keypoint-detection", + "model_class": "VitPoseForPoseEstimation", + "model_type": "vitpose" + }, + "eval": { + "task": "keypoint-detection", + "dataset": { + "path": "~/.cache/winml/datasets/coco_keypoints_val2017", + "build_script": "scripts/build_coco_keypoints.py", + "split": "validation", + "samples": 100, + "shuffle": true, + "seed": 42, + "columns_mapping": { + "input_column": "image", + "annotation_column": "objects", + "keypoints_key": "keypoints", + "bbox_key": "bbox", + "area_key": "area", + "box_format": "xywh" + } + } + } +} diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index b3a6fe3a3..e35a3e35c 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -64,10 +64,33 @@ def convert_to_fp16( if op_block_list: logger.info(" Keeping ops in FP32: %s", op_block_list) + # ORT's converter runs ONNX shape inference by default, which serializes the + # model to a single protobuf message. That serialization fails for models + # whose in-memory proto exceeds protobuf's hard 2 GB message limit + # (google.protobuf.message.EncodeError: Failed to serialize proto), e.g. + # ViT-Huge backbones. Disable shape inference for such models — the FP16 + # cast itself does not require inferred shapes. + protobuf_limit = 2 * 1024**3 + shape_infer_margin = int(protobuf_limit * 0.9) # stay safely under 2 GB + try: + # ByteSize() itself raises once the proto exceeds the 2 GB limit; that + # is already a definitive "too large" signal. + approx_size = model.ByteSize() + except Exception: + approx_size = protobuf_limit + disable_shape_infer = approx_size >= shape_infer_margin + if disable_shape_infer: + logger.info( + " Model is large (~%.2f GB); disabling shape inference during FP16 " + "conversion to avoid the 2 GB protobuf serialization limit.", + approx_size / 1024**3, + ) + converted: ModelProto = convert_float_to_float16( model, keep_io_types=keep_io_types, op_block_list=op_block_list, + disable_shape_infer=disable_shape_infer, ) # ORT's converter appends Cast nodes at the end of the node list (for