search

openvinotoolkit / nncf

Neural Network Compression Framework for enhanced OpenVINOâ„¢ inference
Apache License 2.0
917 stars 228 forks source link

[Torch FX] Support for ExportedProgram from torch.export.export in Weights Compression #2987

Open anzr299 opened 5 days ago

anzr299 commented 5 days ago

🚀 Feature request

Currently, the Torch FX backend supports FX representation(torch.fx.GraphModule() object) extracted using capture_pre_autograd_graph() which has a slightly different representation and is a deprecated feature. This feature request is to compress weights in torch FX graphs extracted using torch.export.export() instead of capture_pre_autograd_graph()

Feature Use Case

The FX model returned from torch.export.export() can also be supported in weights compression along with capture_pre_autograd_graph() due to the latter being deprecated in the future.

Are you going to submit a PR?

daniil-lyakhov commented 5 days ago

@anzr299, let's migrate to torch.export.export(...).module() first. It returns GraphModule, is supported for the target model scope and improves metrics of the swin_v2_s model. #2990