Dumps out module to ``folder`` with ``module_name`` so that it can be imported with ``from import `` Args: folder (Union[str, os.PathLike]): The folder to write the code out to module_name (str): Top-level name to use for the ``Module`
(self, folder: Union[str, os.PathLike], module_name: str = "FxModule")
| 473 | |
| 474 | @compatibility(is_backward_compatible=False) |
| 475 | def to_folder(self, folder: Union[str, os.PathLike], module_name: str = "FxModule"): |
| 476 | """Dumps out module to ``folder`` with ``module_name`` so that it can be |
| 477 | imported with ``from <folder> import <module_name>`` |
| 478 | |
| 479 | Args: |
| 480 | |
| 481 | folder (Union[str, os.PathLike]): The folder to write the code out to |
| 482 | |
| 483 | module_name (str): Top-level name to use for the ``Module`` while |
| 484 | writing out the code |
| 485 | """ |
| 486 | folder = Path(folder) |
| 487 | Path(folder).mkdir(exist_ok=True) |
| 488 | torch.save(self.state_dict(), folder / "state_dict.pt") |
| 489 | tab = " " * 4 |
| 490 | custom_builtins = "\n".join([v.import_str for v in _custom_builtins.values()]) |
| 491 | model_str = f""" |
| 492 | import torch |
| 493 | {custom_builtins} |
| 494 | |
| 495 | from torch.nn import * |
| 496 | class {module_name}(torch.nn.Module): |
| 497 | def __init__(self): |
| 498 | super().__init__() |
| 499 | """ |
| 500 | |
| 501 | def _gen_model_repr(module_name: str, module: torch.nn.Module) -> Optional[str]: |
| 502 | safe_reprs = [ |
| 503 | nn.Linear, |
| 504 | nn.Conv1d, |
| 505 | nn.Conv2d, |
| 506 | nn.Conv3d, |
| 507 | nn.BatchNorm1d, |
| 508 | nn.BatchNorm2d, |
| 509 | nn.BatchNorm3d, |
| 510 | ] |
| 511 | if type(module) in safe_reprs: |
| 512 | return f"{module.__repr__()}" |
| 513 | else: |
| 514 | return None |
| 515 | |
| 516 | blobified_modules = [] |
| 517 | for module_name, module in self.named_children(): |
| 518 | module_str = _gen_model_repr(module_name, module) |
| 519 | if module_str is None: |
| 520 | module_file = folder / f"{module_name}.pt" |
| 521 | torch.save(module, module_file) |
| 522 | blobified_modules.append(module_name) |
| 523 | module_repr = module.__repr__().replace("\r", " ").replace("\n", " ") |
| 524 | module_str = f"torch.load(r'{module_file}') # {module_repr}" |
| 525 | model_str += f"{tab*2}self.{module_name} = {module_str}\n" |
| 526 | |
| 527 | for buffer_name, buffer in self._buffers.items(): |
| 528 | if buffer is None: |
| 529 | continue |
| 530 | model_str += f"{tab*2}self.register_buffer('{buffer_name}', torch.empty({list(buffer.shape)}, dtype={buffer.dtype}))\n" |
| 531 | |
| 532 | for param_name, param in self._parameters.items(): |