[acc_utils] Add print_model_info (#65045)
authorJordan Fix <jfix@fb.com>
Fri, 17 Sep 2021 02:55:46 +0000 (19:55 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 17 Sep 2021 03:29:22 +0000 (20:29 -0700)
commit64641eaee685c858068fee9e78c88c862886854b
tree467081a441508aaa312ef8a85ab50b641d5f7102
parent8c38d141df429459ea6891847950ce157ac82b2c
[acc_utils] Add print_model_info (#65045)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65045

This is a useful tool for printing out all of the ops that are found in a model after acc_tracer. It assumes the provided model has no `call_module` or `call_method`, which is generally a reasonable assumption assuming a model has been successfully traced by the acc_tracer.

Test Plan:
Tested locally. Sample output:
```
Model Info:
> placeholder: 1184
> get_attr: 655
> output: 2
> torch.fx.experimental.fx_acc.acc_ops.add: 2
> torch.fx.experimental.fx_acc.acc_ops.cat: 23
> torch.fx.experimental.fx_acc.acc_ops.embedding_bag: 576
> torch.fx.experimental.fx_acc.acc_ops.layer_norm: 15
> torch.fx.experimental.fx_acc.acc_ops.linear: 27
> torch.fx.experimental.fx_acc.acc_ops.matmul: 3
> torch.fx.experimental.fx_acc.acc_ops.mul: 17
> torch.fx.experimental.fx_acc.acc_ops.permute: 2
> torch.fx.experimental.fx_acc.acc_ops.reshape: 419
> torch.fx.experimental.fx_acc.acc_ops.sigmoid: 16
> torch.fx.experimental.fx_acc.acc_ops.slice_tensor: 630
> torch.fx.experimental.fx_acc.acc_ops.sum: 4
> torch.fx.experimental.fx_acc.acc_ops.tanh: 315
```

Reviewed By: 842974287

Differential Revision: D30954829

fbshipit-source-id: 5c4f0770667b72859b74099d9f4575284fc48bd2
torch/fx/experimental/fx_acc/acc_utils.py