mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
cd28ff5395
Summary:
Adds support for `__getstate__` and `__setstate__` on modules that are called as part of export (`torch.save()`) and import (`torch.jit.load`).
* `__getstate__` and `__setstate__` must be TorchScript functions with the signatures `() -> T` and `(T) -> None` respectively
* The results of `__getstate__` are stored using the pickler in `states.pkl` with one for each module in definition order (`__getstate__` returns `None` by default if an imlpementation is not provided)
* This prevents sharing between `__getstate__` and attributes, but this should be fine since their use is mostly unrelated (attributes are for storing values to be used in script methods, `__getstate__` for running arbitrary computations during import)
Follow up
* Somehow replacing `__getstate__`/`__setstate__` with a `ScriptMethodStub` makes `MyScriptModule().__getstate__()` call `ScriptModule.__getstate__()` when used in Python. This should be fixed so semantics in Python are preserved, but it doesn't affect the typical usage.
](https://our.intern.facebook.com/intern/diff/15287161/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20242
Pulled By: driazati
Differential Revision: D15287161
fbshipit-source-id: b3f5f33ab74a21a89e6d15460af63aff75cab2d8
|
||
|---|---|---|
| .. | ||
| builtin_functions.cpp | ||
| builtin_functions.h | ||
| class_type.cpp | ||
| compilation_unit.h | ||
| compiler.cpp | ||
| compiler.h | ||
| edit_distance.cpp | ||
| edit_distance.h | ||
| error_report.h | ||
| final_returns.cpp | ||
| final_returns.h | ||
| function_schema_parser.cpp | ||
| function_schema_parser.h | ||
| init.cpp | ||
| init.h | ||
| jit_exception.cpp | ||
| jit_exception.h | ||
| lexer.cpp | ||
| lexer.h | ||
| logging.cpp | ||
| logging.h | ||
| module_python.h | ||
| module.cpp | ||
| module.h | ||
| parse_string_literal.h | ||
| parser.cpp | ||
| parser.h | ||
| python_sugared_value.cpp | ||
| python_sugared_value.h | ||
| python_tree_views.cpp | ||
| python_tree_views.h | ||
| resolver.h | ||
| schema_matching.cpp | ||
| schema_matching.h | ||
| schema_type_parser.cpp | ||
| schema_type_parser.h | ||
| script_type_parser.cpp | ||
| script_type_parser.h | ||
| slot.h | ||
| strtod.cpp | ||
| strtod.h | ||
| sugared_value.cpp | ||
| sugared_value.h | ||
| tree_views.h | ||
| tree.h | ||