Skip to content

Commit 2df50c1

Browse files
[TorchFX] Documetation update (#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 #2766
1 parent f0873ca commit 2df50c1

File tree

3 files changed

+47
-10
lines changed

3 files changed

+47
-10
lines changed

README.md

+43-6
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
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.
2121

22-
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).
22+
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).
2323

2424
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).
2525

@@ -33,11 +33,11 @@ learning frameworks.
3333

3434
### Post-Training Compression Algorithms
3535

36-
| Compression algorithm | OpenVINO | PyTorch | TensorFlow | ONNX |
37-
| :------------------------------------------------------------------------------------------------------- | :-------: | :-------: | :-----------: | :-----------: |
38-
| [Post-Training Quantization](./docs/usage/post_training_compression/post_training_quantization/Usage.md) | Supported | Supported | Supported | Supported |
39-
| [Weights Compression](./docs/usage/post_training_compression/weights_compression/Usage.md) | Supported | Supported | Not supported | Not supported |
40-
| [Activation Sparsity](./nncf/experimental/torch/sparsify_activations/ActivationSparsity.md) | Not supported | Experimental |Not supported| Not supported |
36+
| Compression algorithm | OpenVINO | PyTorch | TorchFX | TensorFlow | ONNX |
37+
| :------------------------------------------------------------------------------------------------------- | :-------: | :-------: | :-----------: | :-----------: | :-----------: |
38+
| [Post-Training Quantization](./docs/usage/post_training_compression/post_training_quantization/Usage.md) | Supported | Supported | Experimental | Supported | Supported |
39+
| [Weights Compression](./docs/usage/post_training_compression/weights_compression/Usage.md) | Supported | Supported | Experimental | Not supported | Not supported |
40+
| [Activation Sparsity](./nncf/experimental/torch/sparsify_activations/ActivationSparsity.md) | Not supported | Experimental | Not supported| Not supported| Not supported |
4141

4242
### Training-Time Compression Algorithms
4343

@@ -138,6 +138,42 @@ quantized_model = nncf.quantize(model, calibration_dataset)
138138

139139
</details>
140140

141+
<details><summary><b>TorchFX</b></summary>
142+
143+
```python
144+
import nncf
145+
import torch.fx
146+
from torchvision import datasets, models
147+
from nncf.torch import disable_patching
148+
149+
# Instantiate your uncompressed model
150+
model = models.mobilenet_v2()
151+
152+
# Provide validation part of the dataset to collect statistics needed for the compression algorithm
153+
val_dataset = datasets.ImageFolder("/path", transform=transforms.Compose([transforms.ToTensor()]))
154+
dataset_loader = torch.utils.data.DataLoader(val_dataset)
155+
156+
# Step 1: Initialize the transformation function
157+
def transform_fn(data_item):
158+
images, _ = data_item
159+
return images
160+
161+
# Step 2: Initialize NNCF Dataset
162+
calibration_dataset = nncf.Dataset(dataset_loader, transform_fn)
163+
164+
# Step 3: Export model to TorchFX
165+
input_shape = (1, 3, 224, 224)
166+
with nncf.torch.disable_patching():
167+
fx_model = torch.export.export_for_training(model, args=(ex_input,)).module()
168+
# or
169+
# fx_model = torch.export.export(model, args=(ex_input,)).module()
170+
171+
# Step 4: Run the quantization pipeline
172+
quantized_fx_model = nncf.quantize(fx_model, calibration_dataset)
173+
174+
```
175+
176+
</details>
141177
<details><summary><b>TensorFlow</b></summary>
142178

143179
```python
@@ -384,6 +420,7 @@ Compact scripts demonstrating quantization and corresponding inference speed boo
384420
| [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 |
385421
| [PyTorch MobileNetV2](./examples/post_training_quantization/torch/mobilenet_v2/README.md) | Post-Training Quantization | PyTorch | Image Classification |
386422
| [PyTorch SSD](./examples/post_training_quantization/torch/ssd300_vgg16/README.md) | Post-Training Quantization | PyTorch | Object Detection |
423+
| [TorchFX Resnet18](./examples/post_training_quantization/torch_fx/resnet18/README.md) | Post-Training Quantization | TorchFX | Image Classification |
387424
| [TensorFlow MobileNetV2](./examples/post_training_quantization/tensorflow/mobilenet_v2/README.md) | Post-Training Quantization | TensorFlow | Image Classification |
388425
| [ONNX MobileNetV2](./examples/post_training_quantization/onnx/mobilenet_v2/README.md) | Post-Training Quantization | ONNX | Image Classification |
389426

docs/Algorithms.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@
22

33
## Post-training Compression
44

5-
- [Post Training Quantization (PTQ)](./usage/post_training_compression/post_training_quantization/Usage.md) (OpenVINO, PyTorch, ONNX, TensorFlow)
5+
- [Post Training Quantization (PTQ)](./usage/post_training_compression/post_training_quantization/Usage.md) (OpenVINO, PyTorch, TorchFX, ONNX, TensorFlow)
66
- Symmetric and asymmetric quantization modes
77
- Signed and unsigned
88
- Per tensor/per channel
99
- Each backend support export to the OpenVINO format
10-
- [Weights compression](./usage/post_training_compression/weights_compression/Usage.md) (OpenVINO, PyTorch)
10+
- [Weights compression](./usage/post_training_compression/weights_compression/Usage.md) (OpenVINO, PyTorch, TorchFX)
1111
- Symmetric 8 bit compression mode
1212
- Symmetric and asymmetric 4 bit compression mode
1313
- NF4 compression mode

docs/usage/post_training_compression/post_training_quantization/Usage.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ Every backend has its own return value format for the data transformation functi
5151
backend inference framework.
5252
Below are the formats of data transformation function for each supported backend.
5353

54-
<details><summary><b>PyTorch, TensorFlow, OpenVINO</b></summary>
54+
<details><summary><b>PyTorch, TorchFX, TensorFlow, OpenVINO</b></summary>
5555

5656
The return format of the data transformation function is directly the input tensors consumed by the model. \
5757
_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:
8989
</details>
9090

9191
NNCF provides the examples of Post-Training Quantization where you can find the implementation of data transformation
92-
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)
92+
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)
9393

9494
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).
9595

0 commit comments

Comments
 (0)