[TorchFX] Documetation update (openvinotoolkit#2917)

### Changes * Main README.md, Usage.md and post training quantization docs are updated with info about the TorchFX ### Reason for changes * To reflect new experimental features of TorchFX in the docs ### Related tickets openvinotoolkit#2766
nikita-savelyevv · Nov 21, 2024 · 2ab00cb · 2ab00cb
1 parent f6bbe87
commit 2ab00cb
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -19,7 +19,7 @@
 
 Neural Network Compression Framework (NNCF) provides a suite of post-training and training-time algorithms for optimizing inference of neural networks in [OpenVINO&trade;](https://docs.openvino.ai) with a minimal accuracy drop.
 
-NNCF is designed to work with models from [PyTorch](https://pytorch.org/), [TensorFlow](https://www.tensorflow.org/), [ONNX](https://onnx.ai/) and [OpenVINO&trade;](https://docs.openvino.ai).
+NNCF is designed to work with models from [PyTorch](https://pytorch.org/), [TorchFX](https://pytorch.org/docs/stable/fx.html), [TensorFlow](https://www.tensorflow.org/), [ONNX](https://onnx.ai/) and [OpenVINO&trade;](https://docs.openvino.ai).
 
 NNCF provides [samples](#demos-tutorials-and-samples) that demonstrate the usage of compression algorithms for different use cases and models. See compression results achievable with the NNCF-powered samples on the [NNCF Model Zoo page](./docs/ModelZoo.md).
 
@@ -33,11 +33,11 @@ learning frameworks.
 
 ### Post-Training Compression Algorithms
 
-| Compression algorithm                                                                                    | OpenVINO  | PyTorch   | TensorFlow    | ONNX          |
-| :------------------------------------------------------------------------------------------------------- | :-------: | :-------: | :-----------: | :-----------: |
-| [Post-Training Quantization](./docs/usage/post_training_compression/post_training_quantization/Usage.md) | Supported | Supported | Supported     | Supported     |
-| [Weights Compression](./docs/usage/post_training_compression/weights_compression/Usage.md)               | Supported | Supported | Not supported | Not supported |
-| [Activation Sparsity](./nncf/experimental/torch/sparsify_activations/ActivationSparsity.md)              | Not supported | Experimental |Not supported| Not supported |
+| Compression algorithm                                                                                    | OpenVINO  | PyTorch   | TorchFX   | TensorFlow    | ONNX          |
+| :------------------------------------------------------------------------------------------------------- | :-------: | :-------: | :-----------: | :-----------: | :-----------: |
+| [Post-Training Quantization](./docs/usage/post_training_compression/post_training_quantization/Usage.md) | Supported | Supported | Experimental | Supported     | Supported     |
+| [Weights Compression](./docs/usage/post_training_compression/weights_compression/Usage.md)               | Supported | Supported | Experimental | Not supported | Not supported |
+| [Activation Sparsity](./nncf/experimental/torch/sparsify_activations/ActivationSparsity.md)              | Not supported | Experimental | Not supported| Not supported| Not supported |
 
 ### Training-Time Compression Algorithms
 
@@ -138,6 +138,42 @@ quantized_model = nncf.quantize(model, calibration_dataset)
 
 </details>
 
+<details><summary><b>TorchFX</b></summary>
+
+```python
+import nncf
+import torch.fx
+from torchvision import datasets, models
+from nncf.torch import disable_patching
+
+# Instantiate your uncompressed model
+model = models.mobilenet_v2()
+
+# Provide validation part of the dataset to collect statistics needed for the compression algorithm
+val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
+dataset_loader = torch.utils.data.DataLoader(val_dataset)
+
+# Step 1: Initialize the transformation function
+def transform_fn(data_item):
+    images, _ = data_item
+    return images
+
+# Step 2: Initialize NNCF Dataset
+calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
+
+# Step 3: Export model to TorchFX
+input_shape = (1, 3, 224, 224)
+with nncf.torch.disable_patching():
+    fx_model = torch.export.export_for_training(model, args=(ex_input,)).module()
+    # or
+    # fx_model = torch.export.export(model, args=(ex_input,)).module()
+
+    # Step 4: Run the quantization pipeline
+    quantized_fx_model = nncf.quantize(fx_model, calibration_dataset)
+
+ ```
+
+</details>
 <details><summary><b>TensorFlow</b></summary>
 
 ```python
@@ -384,6 +420,7 @@ Compact scripts demonstrating quantization and corresponding inference speed boo
 | [OpenVINO Anomaly Classification](./examples/post_training_quantization/openvino/anomaly_stfpm_quantize_with_accuracy_control/README.md) | Post-Training Quantization with Accuracy Control |  OpenVINO  | Anomaly Classification |
 | [PyTorch MobileNetV2](./examples/post_training_quantization/torch/mobilenet_v2/README.md)                                                |            Post-Training Quantization            |  PyTorch   |  Image Classification  |
 | [PyTorch SSD](./examples/post_training_quantization/torch/ssd300_vgg16/README.md)                                                        |            Post-Training Quantization            |  PyTorch   |    Object Detection    |
+| [TorchFX Resnet18](./examples/post_training_quantization/torch_fx/resnet18/README.md)                                                        |            Post-Training Quantization            |  TorchFX   |    Image Classification    |
 | [TensorFlow MobileNetV2](./examples/post_training_quantization/tensorflow/mobilenet_v2/README.md)                                        |            Post-Training Quantization            | TensorFlow |  Image Classification  |
 | [ONNX MobileNetV2](./examples/post_training_quantization/onnx/mobilenet_v2/README.md)                                                    |            Post-Training Quantization            |    ONNX    |  Image Classification  |
 

diff --git a/docs/Algorithms.md b/docs/Algorithms.md
@@ -2,12 +2,12 @@
 
 ## Post-training Compression
 
-- [Post Training Quantization (PTQ)](./usage/post_training_compression/post_training_quantization/Usage.md) (OpenVINO, PyTorch, ONNX, TensorFlow)
+- [Post Training Quantization (PTQ)](./usage/post_training_compression/post_training_quantization/Usage.md) (OpenVINO, PyTorch, TorchFX, ONNX, TensorFlow)
   - Symmetric and asymmetric quantization modes
   - Signed and unsigned
   - Per tensor/per channel
   - Each backend support export to the OpenVINO format
-- [Weights compression](./usage/post_training_compression/weights_compression/Usage.md) (OpenVINO, PyTorch)
+- [Weights compression](./usage/post_training_compression/weights_compression/Usage.md) (OpenVINO, PyTorch, TorchFX)
   - Symmetric 8 bit compression mode
   - Symmetric and asymmetric 4 bit compression mode
   - NF4 compression mode

diff --git a/docs/usage/post_training_compression/post_training_quantization/Usage.md b/docs/usage/post_training_compression/post_training_quantization/Usage.md
@@ -51,7 +51,7 @@ Every backend has its own return value format for the data transformation functi
 backend inference framework.
 Below are the formats of data transformation function for each supported backend.
 
-<details><summary><b>PyTorch, TensorFlow, OpenVINO</b></summary>
+<details><summary><b>PyTorch, TorchFX, TensorFlow, OpenVINO</b></summary>
 
 The return format of the data transformation function is directly the input tensors consumed by the model. \
 _If you are not sure that your implementation of data transformation function is correct you can validate it by using the
@@ -89,7 +89,7 @@ for data_item in val_loader:
 </details>
 
 NNCF provides the examples of Post-Training Quantization where you can find the implementation of data transformation
-function: [PyTorch](/examples/post_training_quantization/torch/mobilenet_v2/README.md), [TensorFlow](/examples/post_training_quantization/tensorflow/mobilenet_v2/README.md), [ONNX](/examples/post_training_quantization/onnx/mobilenet_v2/README.md), and [OpenVINO](/examples/post_training_quantization/openvino/mobilenet_v2/README.md)
+function: [PyTorch](/examples/post_training_quantization/torch/mobilenet_v2/README.md), [TorchFX](/examples/post_training_quantization/torch_fx/resnet18/README.md), [TensorFlow](/examples/post_training_quantization/tensorflow/mobilenet_v2/README.md), [ONNX](/examples/post_training_quantization/onnx/mobilenet_v2/README.md), and [OpenVINO](/examples/post_training_quantization/openvino/mobilenet_v2/README.md)
 
 In case the Post-Training Quantization algorithm could not reach quality requirements you can fine-tune a quantized pytorch model. Example of the Quantization-Aware training pipeline for a pytorch model could be found [here](/examples/quantization_aware_training/torch/resnet18/README.md).