pytorch

431d25a1 - [export] Add save/load function (#107309)

[export] Add save/load function (#107309) Added the following APIs: ``` def save( ep: ExportedProgram, f: Union[str, pathlib.Path, io.BytesIO], extra_files: Optional[Dict[str, Any]] = None, opset_version: Optional[Dict[str, int]] = None, ) -> None: """ Saves a version of the given exported program for use in a separate process. Args: ep (ExportedProgram): The exported program to save. f (str): A file-like object (has to implement write and flush) or a string containing a file name. extra_files (Optional[Dict[str, Any]]): Map from filename to contents which will be stored as part of f. opset_version (Optional[Dict[str, int]]): A map of opset names to the version of this opset """ def load( f: Union[str, pathlib.Path, io.BytesIO], extra_files: Optional[Dict[str, Any]] = None, expected_opset_version: Optional[Dict[str, int]] = None, ) -> ExportedProgram: """ Loads an ExportedProgram previously saved with torch._export.save Args: ep (ExportedProgram): The exported program to save. f (str): A file-like object (has to implement write and flush) or a string containing a file name. extra_files (Optional[Dict[str, Any]]): The extra filenames given in this map would be loaded and their content would be stored in the provided map. expected_opset_version (Optional[Dict[str, int]]): A map of opset names to expected opset versions Returns: An ExportedProgram object """ ``` Example usage: ``` # With buffer buffer = io.BytesIO() torch._export.save(ep, buffer) buffer.seek(0) loaded_ep = torch._export.load(buffer) # With file with tempfile.NamedTemporaryFile() as f: torch._export.save(ep, f) f.seek(0) loaded_ep = torch._export.load(f) # With Path with TemporaryFileName() as fname: path = pathlib.Path(fname) torch._export.save(ep, path) loaded_ep = torch._export.load(path) # Saving with extra files buffer = io.BytesIO() save_extra_files = {"extra.txt": "moo"} torch._export.save(ep, buffer, save_extra_files) buffer.seek(0) load_extra_files = {"extra.txt": ""} loaded_ep = torch._export.load(buffer, extra_files) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/107309 Approved by: https://github.com/avikchaudhuri, https://github.com/gmagogsfm, https://github.com/tugsbayasgalan