| 51 | |
| 52 | |
| 53 | def parse_args() -> argparse.Namespace: |
| 54 | parser = argparse.ArgumentParser() |
| 55 | |
| 56 | parser.add_argument( |
| 57 | "-m", |
| 58 | "--model_name", |
| 59 | required=True, |
| 60 | help=f"Provide model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}", |
| 61 | ) |
| 62 | |
| 63 | parser.add_argument( |
| 64 | "-c", |
| 65 | "--compute_unit", |
| 66 | required=False, |
| 67 | default=ct.ComputeUnit.ALL.name.lower(), |
| 68 | help=f"Provide compute unit for the model. Valid ones: {[[compute_unit.name.lower() for compute_unit in ct.ComputeUnit]]}", |
| 69 | ) |
| 70 | |
| 71 | parser.add_argument( |
| 72 | "-precision", |
| 73 | "--compute_precision", |
| 74 | required=False, |
| 75 | default=ct.precision.FLOAT16.value, |
| 76 | help=f"Provide compute precision for the model. Valid ones: {[[precision.value for precision in ct.precision]]}", |
| 77 | ) |
| 78 | |
| 79 | parser.add_argument( |
| 80 | "--compile", |
| 81 | action=argparse.BooleanOptionalAction, |
| 82 | required=False, |
| 83 | default=False, |
| 84 | ) |
| 85 | parser.add_argument("--use_partitioner", action=argparse.BooleanOptionalAction) |
| 86 | parser.add_argument("--generate_etrecord", action=argparse.BooleanOptionalAction) |
| 87 | parser.add_argument("--save_processed_bytes", action=argparse.BooleanOptionalAction) |
| 88 | parser.add_argument( |
| 89 | "--dynamic_shapes", |
| 90 | action=argparse.BooleanOptionalAction, |
| 91 | required=False, |
| 92 | default=False, |
| 93 | ) |
| 94 | parser.add_argument( |
| 95 | "--run_with_pybindings", |
| 96 | action=argparse.BooleanOptionalAction, |
| 97 | ) |
| 98 | |
| 99 | args = parser.parse_args() |
| 100 | return args |
| 101 | |
| 102 | |
| 103 | def partition_module_to_coreml(module): |