()
| 212 | |
| 213 | |
| 214 | def main() -> None: |
| 215 | known_backend_ids = [ |
| 216 | BackendWithCompilerDemo.__name__, |
| 217 | StubBackend.__name__, |
| 218 | ] |
| 219 | |
| 220 | # These args are optimized for genrule usage. There's a lot of startup |
| 221 | # overhead for this tool, so it's faster to export multiple models at once |
| 222 | # when possible. |
| 223 | parser = argparse.ArgumentParser( |
| 224 | prog="export_delegated_program", |
| 225 | description="Exports delegated nn.Module models to ExecuTorch .pte files", |
| 226 | ) |
| 227 | parser.add_argument( |
| 228 | "--modules", |
| 229 | help="Comma-separated list of model class names to export; " |
| 230 | + "e.g., '--modules=ModuleOne,ModuleTwo'", |
| 231 | type=lambda s: [item.strip() for item in s.split(",")], |
| 232 | ) |
| 233 | parser.add_argument( |
| 234 | "--backend_id", |
| 235 | type=str, |
| 236 | default=StubBackend.__name__, |
| 237 | help="ID of the backend to use for delegation; " |
| 238 | + f"one of {known_backend_ids}", |
| 239 | ) |
| 240 | parser.add_argument( |
| 241 | "--inline_delegate_segments", |
| 242 | action="store_true", |
| 243 | help="Store delegate data inside the flatbuffer.", |
| 244 | ) |
| 245 | parser.add_argument( |
| 246 | "--delegate_alignment", type=int, default=None, help="Delegate alignment." |
| 247 | ) |
| 248 | parser.add_argument( |
| 249 | "--external_constants", |
| 250 | action="store_true", |
| 251 | help="Export the model with all constants saved to an external file.", |
| 252 | ) |
| 253 | parser.add_argument( |
| 254 | "--outdir", |
| 255 | type=str, |
| 256 | required=True, |
| 257 | help="Path to the directory to write <classname>[-<suffix>[...]].pte " |
| 258 | + "files to.", |
| 259 | ) |
| 260 | args = parser.parse_args() |
| 261 | |
| 262 | # Find the classes to export. Only looks in this module for now, but could |
| 263 | # be extended to look in other modules if helpful. |
| 264 | module_names_to_classes: Dict[str, Type[nn.Module]] = {} |
| 265 | for module in args.modules: |
| 266 | module_class = getattr(sys.modules[__name__], module, None) |
| 267 | if not (inspect.isclass(module_class) and issubclass(module_class, nn.Module)): |
| 268 | raise NameError(f"Could not find nn.Module class named '{module}'") |
| 269 | module_names_to_classes[module] = module_class |
| 270 | |
| 271 | # Export and write to the output files. |
no test coverage detected