From 9d0a87f4e98196259f32851fb69f43198863d824 Mon Sep 17 00:00:00 2001 From: dlyakhov Date: Thu, 14 Nov 2024 14:05:32 +0100 Subject: [PATCH] [TorchFX] Documetation update --- README.md | 47 ++++++++++++++++--- docs/Algorithms.md | 4 +- .../post_training_quantization/Usage.md | 4 +- 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index f4b58c3d710..bbc455e5e2e 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,40 @@ 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() + + # Step 4: Run the quantization pipeline + quantized_fx_model = nncf.quantize(fx_model, calibration_dataset) + + ``` + +
TensorFlow ```python @@ -384,6 +418,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:
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).