ns for fx: rekey results by model node names (#60305)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60305
Adjusts the NS for FX weight and activation extraction APIs
to require a model name, and rekeys the results of these APIs
to use the node names of the specified model as layer keys.
For example, before
```
// API call
results = ns.extract_logger_info(
model_a, model_b, ns.OutputLogger)
// results
{'base_op_1_0': {'node_output':
{'model_a': [{'ref_node_name': 'linear1', ...}]}}}
```
and after
```
// API call
results = ns.extract_logger_info(
model_a, model_b, ns.OutputLogger, 'model_b_name')
// results
// note: instead of `base_op_1_0`, the layer is named `linear1`
{'linear1': {'node_output':
{'model_a': [{'ref_node_name': 'linear1', ...}]}}}
```
Note: we cannot use these names while collecting data because
node names are not guaranteed to be consistent across graphs.
This is why we only rekey as the very last step.
Test Plan:
```
python test/test_quantization.py TestFXNumericSuiteCoreAPIs.test_layer_names
```
Imported from OSS
Reviewed By: hx89
Differential Revision: D29243045
fbshipit-source-id: d39ecdfdd18b07291e3ecefed2ede287b100b7d0
