Skip to content

Commit

Permalink
[TorchFX] Documetation update
Browse files Browse the repository at this point in the history
  • Loading branch information
daniil-lyakhov committed Nov 14, 2024
1 parent 90d15a6 commit 9d0a87f
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 10 deletions.
47 changes: 41 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).

Expand All @@ -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

Expand Down Expand Up @@ -138,6 +138,40 @@ 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()

# Step 4: Run the quantization pipeline
quantized_fx_model = nncf.quantize(fx_model, calibration_dataset)

```

</details>
<details><summary><b>TensorFlow</b></summary>

```python
Expand Down Expand Up @@ -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 |

Expand Down
4 changes: 2 additions & 2 deletions docs/Algorithms.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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).

Expand Down

0 comments on commit 9d0a87f

Please sign in to comment.