diff --git a/README.md b/README.md
index f4b58c3d710..6c263b4c34a 100644
--- 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™](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™](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™](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)
+TorchFX
+
+```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)
+
+ ```
+
+TensorFlow
```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
index d6e23cebfdb..992f8183962 100644
--- 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
index 0a078a4399f..c98fe38f298 100644
--- 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.
-PyTorch, TensorFlow, OpenVINO
+PyTorch, TorchFX, TensorFlow, OpenVINO
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: