pytorch
5ad3bc71 - ns for fx: change node I/O determination to strict allowlist (#55434)

Commit
3 years ago
ns for fx: change node I/O determination to strict allowlist (#55434) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55434 Before this PR, there was some hacky logic which determined the input and output types of nodes based on heuristics such as inspecting `__module__`, or assuming that an op has an I/O dtype of `torch.float` when the heuristics did not find any matches. This is problematic because the heuristics were not exact, and this could result in non-sensical shadow graphs when the heuristics would return an incorrect dtype. This PR switches the dtype determination to an allowlist system, where we specify exactly what the dtypes are for the nodes or modules which are in an allowlist, and we add an `UNKNOWN` type for everything else. The shadow logic is changed to skip inserting shadows on any function or module where the I/O dtype is unknown. The current allowlist only contains functions necessary for the currently existing tests. Filling out the allowlist with all necessary torch functions is left for a future PR. As a result of this, we can do the following (also implemented in this PR): 1. enable graph matching on nodes with equal types (for example, F.linear and F.linear). The restriction that only nodes with equal types was in the code as a placeholder, it's better to allow comparisons of nodes of equal types. One case where this is useful is unshadowed activations. 2. enable models with user defined modules to be passed to Numeric Suite APIs without errors. Test Plan: ``` python test/test_quantization.py TestFXGraphMatcher python test/test_quantization.py TestFXGraphMatcherModels python test/test_quantization.py TestFXNumericSuiteCoreAPIs python test/test_quantization.py TestFXNumericSuiteCoreAPIsModels ``` Imported from OSS Reviewed By: jerryzh168 Differential Revision: D27622418 fbshipit-source-id: 40dcba0222c01154c141467640c1eb89725f33a7
Author
Parents
Loading