()
| 198 | |
| 199 | |
| 200 | def main(): |
| 201 | args = parse_args() |
| 202 | |
| 203 | if args.model_name not in MODEL_NAME_TO_MODEL: |
| 204 | raise RuntimeError( |
| 205 | f"Model {args.model_name} is not a valid name. " |
| 206 | f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}." |
| 207 | ) |
| 208 | |
| 209 | valid_compute_units = [compute_unit.name.lower() for compute_unit in ct.ComputeUnit] |
| 210 | if args.compute_unit not in valid_compute_units: |
| 211 | raise RuntimeError( |
| 212 | f"{args.compute_unit} is invalid. " |
| 213 | f"Valid compute units are {valid_compute_units}." |
| 214 | ) |
| 215 | |
| 216 | model, example_args, example_kwargs, dynamic_shapes = ( |
| 217 | EagerModelFactory.create_model(*MODEL_NAME_TO_MODEL[args.model_name]) |
| 218 | ) |
| 219 | if not args.dynamic_shapes: |
| 220 | dynamic_shapes = None |
| 221 | |
| 222 | compile_specs = generate_compile_specs_from_args(args) |
| 223 | pte_base_name = get_pte_base_name(args) |
| 224 | if args.use_partitioner: |
| 225 | model = model.eval() |
| 226 | ep = torch.export.export( |
| 227 | model, |
| 228 | args=example_args, |
| 229 | kwargs=example_kwargs, |
| 230 | dynamic_shapes=dynamic_shapes, |
| 231 | ) |
| 232 | print(ep) |
| 233 | delegated_program = exir.to_edge_transform_and_lower( |
| 234 | ep, |
| 235 | partitioner=[CoreMLPartitioner(compile_specs=compile_specs)], |
| 236 | generate_etrecord=args.generate_etrecord, |
| 237 | ) |
| 238 | exec_program = delegated_program.to_executorch() |
| 239 | save_pte_program(exec_program, pte_base_name) |
| 240 | if args.generate_etrecord: |
| 241 | exec_program.get_etrecord().save(f"{pte_base_name}_coreml_etrecord.bin") |
| 242 | if args.run_with_pybindings: |
| 243 | run_with_pybindings( |
| 244 | executorch_program=exec_program, |
| 245 | eager_reference=model, |
| 246 | example_inputs=example_args, |
| 247 | precision=args.compute_precision, |
| 248 | ) |
| 249 | else: |
| 250 | lowered_module, edge_copy = lower_module_to_coreml( |
| 251 | module=model, |
| 252 | example_inputs=example_args, |
| 253 | compile_specs=compile_specs, |
| 254 | ) |
| 255 | exec_program = export_lowered_module_to_executorch_program( |
| 256 | lowered_module, |
| 257 | example_args, |
no test coverage detected