diff --git a/README.md b/README.md
index 4347feec8cf..65bbb7f3e5c 100644
--- a/README.md
+++ b/README.md
@@ -201,6 +201,9 @@ def transform_fn(data_item):
calibration_dataset = nncf.Dataset(val_dataset, transform_fn)
# Step 3: Run the quantization pipeline
quantized_model = nncf.quantize(model, calibration_dataset)
+# Step 4: Remove auxiliary layers and operations added during the quantization process,
+# resulting in a clean, fully quantized model ready for deployment.
+stripped_model = nncf.strip(quantized_model)
```
diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 3a5fbffb096..8ae7163e677 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -1,7 +1,7 @@
-# Use NNCF for Quantization Aware Training in PyTorch
+# Use NNCF for Quantization Aware Training
-This is a step-by-step tutorial on how to integrate the NNCF package into the existing PyTorch project (please see the [TensorFlow quantization documentation](../other_algorithms/LegacyQuantization.md) for integration tutorial for the existing TensorFlow project).
-The use case implies that the user already has a training pipeline that reproduces training of the model in the floating point precision and pretrained model.
+This is a step-by-step tutorial on how to integrate the NNCF package into the existing PyTorch or TensorFlow projects.
+The use case implies that the user already has a training pipeline that reproduces training of the model in the floating point precision and pretrained model.
The task is to prepare this model for accelerated inference by simulating the compression at train time.
Please refer to this [document](/docs/usage/training_time_compression/other_algorithms/LegacyQuantization.md) for details of the implementation.
@@ -11,11 +11,24 @@ Please refer to this [document](/docs/usage/training_time_compression/other_algo
Quantize the model using the [Post Training Quantization](../../post_training_compression/post_training_quantization/Usage.md) method.
+PyTorch
+
```python
model = TorchModel() # instance of torch.nn.Module
quantized_model = nncf.quantize(model, ...)
```
+TensorFlow
+
+```python
+model = TensorFlowModel() # instance of tf.keras.Model
+quantized_model = nncf.quantize(model, ...)
+```
+
+PyTorch
+
+Trace the model via inference in framework operations.
-1. Trace the model via inference in framework operations.
+```python
+# To OpenVINO format
+import openvino as ov
+ov_quantized_model = ov.convert_model(quantized_model.cpu(), example_input=dummy_input)
+```
+
+TensorFlow
+
+```python
+# To OpenVINO format
+import openvino as ov
+
+# Removes auxiliary layers and operations added during the quantization process,
+# resulting in a clean, fully quantized model ready for deployment.
+stripped_model = nncf.strip(quantized_model)
+
+ov_quantized_model = ov.convert_model(stripped_model)
+```
- ```python
- # To OpenVINO format
- import openvino as ov
- ov_quantized_model = ov.convert_model(quantized_model.cpu(), example_input=dummy_input)
- ```
+PyTorch
+
The complete information about compression is defined by a compressed model and a NNCF config.
The model characterizes the weights and topology of the network. The NNCF config - how to restore additional modules intoduced by NNCF.
The NNCF config can be obtained by `quantized_model.nncf.get_config()` on saving and passed to the
@@ -46,8 +80,6 @@ The NNCF config can be obtained by `quantized_model.nncf.get_config()` on saving
The quantized model saving allows to load quantized modules to the target model in a new python process and
requires only example input for the target module, corresponding NNCF config and the quantized model state dict.
-### Saving and loading compressed models in PyTorch
-
```python
# save part
quantized_model = nncf.quantize(model, calibration_dataset)
@@ -70,10 +102,53 @@ quantized_model.load_state_dict(state_dict)
You can save the `compressed_model` object `torch.save` as usual: via `state_dict` and `load_state_dict` methods.
+TensorFlow
+
+To save a model checkpoint, use the following API:
+
+```python
+from nncf.tensorflow import ConfigState
+from nncf.tensorflow import get_config
+from nncf.tensorflow.callbacks.checkpoint_callback import CheckpointManagerCallback
+
+nncf_config = get_config(quantized_model)
+checkpoint = tf.train.Checkpoint(model=quantized_model,
+ nncf_config_state=ConfigState(nncf_config),
+ ... # the rest of the user-defined objects to save
+ )
+callbacks = []
+callbacks.append(CheckpointManagerCallback(checkpoint, path_to_checkpoint))
+...
+quantized_model.fit(..., callbacks=callbacks)
+```
+
+To restore the model from checkpoint, use the following API:
+
+```python
+from nncf.tensorflow import ConfigState
+from nncf.tensorflow import load_from_config
+
+checkpoint = tf.train.Checkpoint(nncf_config_state=ConfigState())
+checkpoint.restore(path_to_checkpoint)
+
+quantized_model = load_from_config(model, checkpoint.nncf_config_state.config)
+
+checkpoint = tf.train.Checkpoint(model=quantized_model
+ ... # the rest of the user-defined objects to load
+ )
+checkpoint.restore(path_to_checkpoint)
+```
+
+PyTorch
+
With no target model code modifications, NNCF only supports native PyTorch modules with respect to trainable parameter (weight) compressed, such as `torch.nn.Conv2d`.
If your model contains a custom, non-PyTorch standard module with trainable weights that should be compressed, you can register it using the `@nncf.register_module` decorator:
@@ -91,4 +166,9 @@ If registered module should be ignored by specific algorithms use `ignored_algor
In the example above, the NNCF-compressed models that contain instances of `MyModule` will have the corresponding modules extended with functionality that will allow NNCF to quantize the `weight` parameter of `MyModule` before it takes part in `MyModule`'s `forward` calculation.
-See a PyTorch [example](/examples/quantization_aware_training/torch/resnet18/README.md) for **Quantization** Compression scenario on Tiny ImageNet-200 dataset.
+