OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 37 Packages Projects Releases Wiki Activity
12daa4f663
pytorch/torch/jit
History
Mike Iovine dd81fa9027 [JIT] Freeze allows preservation of submodule attributes (#66102) 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66102

This changes allows the `preserved_attributes` parameter of `torch.jit.freeze` to accept attributes of submodules. Previously, only root-level attributes were able to be preserved. Example:

```
class SubModule(nn.Module):
    def __init__(self):
        super(SubModule, self).__init__()
        self.a = 1
        self.b = 2

    def forward(self):
        return self.a + self.b

class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.sub = SubModule()

    def forward(self):
        return self.sub()

mod = torch.jit.script(Module())
mod.eval()
frozen_mod = torch.jit.freeze(mod, preserved_attrs = ['sub.a'])

mod.sub   # OK
mod.sub.a # OK
mod.sub.b # Error, not preserved
mod()     # = 3
mod.sub.a = 0
mod()     # = 2
```

Test Plan: `buck test caffe2/test:jit -- TestFreezing`

Reviewed By: eellison

Differential Revision: D31383868

fbshipit-source-id: 34a05ca9528d4e5f04f71ac2a339fd584a8fa305
2021-10-25 07:56:20 -07:00
..
mobile [PyTorch Edge][type] Add type check in compatibility api (#63129) 2021-10-06 02:23:44 -07:00
__init__.py Back out "D30740897 Add fusion enabled apis" (#64500) 2021-09-04 20:55:58 -07:00
_async.py [docs] Fix backticks in docs (#60474) 2021-06-24 06:27:41 -07:00
_builtins.py [TS] Add complex support for more ops (#54541) 2021-05-12 23:33:29 -07:00
_check.py Add lint for unqualified type: ignore (#56290) 2021-04-21 08:07:23 -07:00
_freeze.py [JIT] Freeze allows preservation of submodule attributes (#66102) 2021-10-25 07:56:20 -07:00
_fuser.py patching graph_for (#55139) 2021-10-04 21:50:22 -07:00
_logging.py
_monkeytype_config.py More robust check of whether a class is defined in torch (#64083) 2021-08-27 08:55:35 -07:00
_pickle.py [jit] Add type tags to lists/dicts in pickle (#33255) 2020-03-10 19:17:01 -07:00
_recursive.py clean up mypy nit in torch/jit/_recursive.py (#66253) 2021-10-08 18:07:33 -07:00
_script.py patching graph_for (#55139) 2021-10-04 21:50:22 -07:00
_serialization.py Add lint for unqualified type: ignore (#56290) 2021-04-21 08:07:23 -07:00
_state.py [package] fix mangling issues with TorchScript (#54915) 2021-03-31 00:58:05 -07:00
_trace.py Updates internal assert_allclose callsites in favor of assert_close (#61841) 2021-08-19 12:50:41 -07:00
annotations.py Support Union in TorchScript (#64234) 2021-09-03 06:12:24 -07:00
frontend.py changes for pytorch issue 55577 (#66571) 2021-10-14 18:19:11 -07:00
quantized.py pytorch quantization ao migration phase 2: torch/jit (#65829) 2021-10-01 06:22:22 -07:00
supported_ops.py Add lint for unqualified noqa (#56272) 2021-04-19 13:16:18 -07:00
unsupported_tensor_ops.py