Spaces:
Configuration error
Configuration error
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| # | |
| # Copyright (c) 2022 Intel Corporation | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| import click | |
| import inspect | |
| import os | |
| import sys | |
| from tlt.utils.types import FrameworkType | |
| from tlt.utils.inc_utils import get_inc_config | |
| def quantize(model_dir, dataset_dir, dataset_file, delimiter, batch_size, approach, accuracy_criterion, timeout, | |
| max_trials, output_dir): | |
| """ | |
| Uses the Intel Neural Compressor to perform post-training quantization on a trained model | |
| """ | |
| print("Model directory:", model_dir) | |
| print("Dataset directory:", dataset_dir) | |
| print("Quantization approach:", approach) | |
| print("Accuracy criterion:", accuracy_criterion) | |
| print("Exit policy timeout:", timeout) | |
| print("Exit policy max trials:", max_trials) | |
| print("Batch size:", batch_size) | |
| print("Output directory:", output_dir) | |
| try: | |
| # Create the output directory, if it doesn't exist | |
| from tlt.utils.file_utils import verify_directory | |
| verify_directory(output_dir, require_directory_exists=False) | |
| except Exception as e: | |
| sys.exit("Error while verifying the output directory: {}", str(e)) | |
| saved_model_path = os.path.join(model_dir, "saved_model.pb") | |
| pytorch_model_path = os.path.join(model_dir, "model.pt") | |
| if os.path.isfile(saved_model_path): | |
| framework = FrameworkType.TENSORFLOW | |
| elif os.path.isfile(pytorch_model_path): | |
| framework = FrameworkType.PYTORCH | |
| else: | |
| sys.exit("Quantization is currently only implemented for TensorFlow saved_model.pb and PyTorch model.pt " | |
| "models. No such files found in the model directory ({}).".format(model_dir)) | |
| # Get the model name from the directory path, assuming models are exported like <model name>/n | |
| model_name = os.path.basename(os.path.dirname(model_dir)) | |
| print("Model name:", model_name) | |
| print("Framework:", framework) | |
| try: | |
| from tlt.models.model_factory import get_model | |
| model = get_model(model_name, framework) | |
| model.load_from_directory(model_dir) | |
| except Exception as e: | |
| sys.exit("An error occurred while getting the model: {}\nNote that the model directory is expected to contain " | |
| "a previously exported model where the directory structure is <model name>/n/saved_model.pb " | |
| "(for TensorFlow) or <model name>/n/model.pt (for PyTorch).".format(str(e))) | |
| try: | |
| from tlt.datasets import dataset_factory | |
| if str(model.use_case) == "image_classification": | |
| dataset = dataset_factory.load_dataset(dataset_dir, model.use_case, model.framework) | |
| elif str(model.use_case) == "text_classification": | |
| if not dataset_file: | |
| raise ValueError("Loading a text classification dataset requires --dataset-file to specify the " | |
| "file name of the .csv file to load from the --dataset-dir.") | |
| if not delimiter: | |
| raise ValueError("Loading a text classification dataset requires --delimiter in order to read the " | |
| ".csv file from the --dataset-dir. in the correct format") | |
| dataset = dataset_factory.load_dataset(dataset_dir, model.use_case, model.framework, | |
| csv_file_name=dataset_file, delimiter=delimiter) | |
| else: | |
| sys.exit("ERROR: Quantization is currently only implemented for Image Classification " | |
| "and Text Classification models") | |
| # Preprocess, batch, and split | |
| if 'image_size' in inspect.getfullargspec(dataset.preprocess).args: # For Image classification | |
| dataset.preprocess(image_size=model.image_size, batch_size=batch_size) | |
| elif 'model_name' in inspect.getfullargspec(dataset.preprocess).args: # For HF Text classification | |
| dataset.preprocess(model_name=model_name, batch_size=batch_size) | |
| else: # For TF Text classification | |
| dataset.preprocess(batch_size=batch_size) | |
| dataset.shuffle_split() | |
| # Generate a default inc config | |
| inc_config = get_inc_config(approach, accuracy_criterion, timeout, max_trials) | |
| # Setup a directory for the quantized model | |
| quantized_output_dir = os.path.join(output_dir, "quantized", model_name) | |
| verify_directory(quantized_output_dir) | |
| if len(os.listdir(quantized_output_dir)) > 0: | |
| quantized_output_dir = os.path.join(quantized_output_dir, "{}".format( | |
| len(os.listdir(quantized_output_dir)) + 1)) | |
| else: | |
| quantized_output_dir = os.path.join(quantized_output_dir, "1") | |
| # Call the quantization API | |
| print("Starting post-training quantization", flush=True) | |
| model.quantize(quantized_output_dir, dataset, config=inc_config) | |
| except Exception as e: | |
| sys.exit("An error occurred during quantization: {}".format(str(e))) | |