()
| 355 | |
| 356 | |
| 357 | def main() -> None: |
| 358 | # These args are optimized for genrule usage. There's a lot of startup |
| 359 | # overhead for this tool, so it's faster to export multiple models at once |
| 360 | # when possible. |
| 361 | torch.manual_seed(0) |
| 362 | parser = argparse.ArgumentParser( |
| 363 | prog="export_program", |
| 364 | description="Exports nn.Module models to ExecuTorch .pte files", |
| 365 | ) |
| 366 | parser.add_argument( |
| 367 | "--modules", |
| 368 | help="Comma-separated list of model class names to export; " |
| 369 | + "e.g., '--modules=ModuleBasic,ModuleAdd'", |
| 370 | type=lambda s: [item.strip() for item in s.split(",")], |
| 371 | ) |
| 372 | parser.add_argument( |
| 373 | "--outdir", |
| 374 | type=str, |
| 375 | required=True, |
| 376 | help="Path to the directory to write <classname>.pte files and .ptd files to", |
| 377 | ) |
| 378 | parser.add_argument( |
| 379 | "--external-constants", |
| 380 | action="store_true", |
| 381 | help="Export the model with external constants", |
| 382 | ) |
| 383 | args = parser.parse_args() |
| 384 | |
| 385 | # Find the classes to export. Only looks in this module for now, but could |
| 386 | # be extended to look in other modules if helpful. |
| 387 | module_names_to_classes: Dict[str, Type[nn.Module]] = {} |
| 388 | for module in args.modules: |
| 389 | module_class = getattr(sys.modules[__name__], module, None) |
| 390 | if not (inspect.isclass(module_class) and issubclass(module_class, nn.Module)): |
| 391 | raise NameError(f"Could not find nn.Module class named '{module}'") |
| 392 | module_names_to_classes[module] = module_class |
| 393 | |
| 394 | # Export and write to the output files. |
| 395 | os.makedirs(args.outdir, exist_ok=True) |
| 396 | for module_name, module_class in module_names_to_classes.items(): |
| 397 | if args.external_constants: |
| 398 | module_name = f"{module_name}Program" |
| 399 | outfile = os.path.join(args.outdir, f"{module_name}.pte") |
| 400 | prog = export_module_to_program( |
| 401 | module_class, |
| 402 | external_constants=args.external_constants, |
| 403 | ) |
| 404 | with open(outfile, "wb") as fp: |
| 405 | prog.write_to_file(fp) |
| 406 | print(f"Exported {module_name} and wrote program data to {outfile}") |
| 407 | |
| 408 | if args.external_constants: |
| 409 | # current infra doesnt easily allow renaming this file, so just hackily do it here. |
| 410 | prog._tensor_data[f"{module_name}"] = prog._tensor_data.pop( |
| 411 | "_default_external_constant" |
| 412 | ) |
| 413 | prog.write_tensor_data_to_file(args.outdir) |
| 414 |
no test coverage detected