| 1427 | |
| 1428 | |
| 1429 | class GraphModuleDeserializer: |
| 1430 | @dataclasses.dataclass |
| 1431 | class Result: |
| 1432 | graph_module: torch.fx.GraphModule |
| 1433 | signature: ep.ExportGraphSignature |
| 1434 | module_call_graph: List[ep.ModuleCallEntry] |
| 1435 | names_to_symbols: Dict[str, sympy.Symbol] |
| 1436 | state_dict: Dict[str, Union[torch.Tensor, torch.nn.Parameter]] |
| 1437 | constants: Dict[str, Union[torch.Tensor, torch.ScriptObject]] |
| 1438 | example_inputs: Optional[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]]] |
| 1439 | |
| 1440 | def __init__(self): |
| 1441 | self.serialized_name_to_node: Dict[str, torch.fx.Node] = {} |
| 1442 | self.serialized_name_to_meta: Dict[str, MetaType] = {} |
| 1443 | self.graph = torch.fx.Graph() |
| 1444 | self.module = torch.nn.Module() |
| 1445 | |
| 1446 | @contextmanager |
| 1447 | def save_graph_module(self) -> Iterator[None]: |
| 1448 | saved = ( |
| 1449 | self.graph, |
| 1450 | self.module, |
| 1451 | self.serialized_name_to_node, |
| 1452 | self.serialized_name_to_meta, |
| 1453 | ) |
| 1454 | self.graph = torch.fx.Graph() |
| 1455 | self.module = torch.nn.Module() |
| 1456 | self.serialized_name_to_node = {} |
| 1457 | self.serialized_name_to_meta = {} |
| 1458 | try: |
| 1459 | yield |
| 1460 | finally: |
| 1461 | ( |
| 1462 | self.graph, |
| 1463 | self.module, |
| 1464 | self.serialized_name_to_node, |
| 1465 | self.serialized_name_to_meta, |
| 1466 | ) = saved |
| 1467 | |
| 1468 | def deserialize_operator(self, serialized_target: str): |
| 1469 | if serialized_target.startswith( |
| 1470 | "_operator" |
| 1471 | ): # TODO(zhxchen17) Follow up on this. |
| 1472 | module = operator |
| 1473 | serialized_target_names = serialized_target.split(".")[1:] |
| 1474 | elif serialized_target.startswith("torch"): |
| 1475 | module = torch # type: ignore[misc] |
| 1476 | serialized_target_names = serialized_target.split(".")[1:] |
| 1477 | else: # TODO(zhxchen17) Don't catch all here. |
| 1478 | return serialized_target |
| 1479 | |
| 1480 | target = module |
| 1481 | for name in serialized_target_names: |
| 1482 | if not hasattr(target, name): |
| 1483 | return serialized_target |
| 1484 | else: |
| 1485 | target = getattr(target, name) |
| 1486 | return target |