Spaces:
Configuration error
Configuration error
File size: 3,725 Bytes
a01ef8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
#!/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 os
import sys
from tlt.utils.types import FrameworkType
@click.command()
@click.option("--model-dir", "--model_dir",
required=True,
type=click.Path(exists=True, file_okay=False),
help="Model directory to reload for graph optimization. The model directory should contain a "
"saved_model.pb TensorFlow model.")
@click.option("--output-dir", "--output_dir",
required=True,
type=click.Path(file_okay=False),
help="A writeable output directory. The output directory will be used as a location to save the "
"optimized model.")
def optimize(model_dir, output_dir):
"""
Uses the Intel Neural Compressor to perform graph optimization on a trained model
"""
print("Model directory:", model_dir)
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
else:
sys.exit("Graph optimization is currently only supported for TensorFlow saved_model.pb "
"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).".format(str(e)))
try:
# Setup a directory for the quantized model
optimized_output_dir = os.path.join(output_dir, "optimized", model_name)
verify_directory(optimized_output_dir)
if len(os.listdir(optimized_output_dir)) > 0:
optimized_output_dir = os.path.join(optimized_output_dir, "{}".format(
len(os.listdir(optimized_output_dir)) + 1))
else:
optimized_output_dir = os.path.join(optimized_output_dir, "1")
# Call the graph optimization API
print("Starting graph optimization", flush=True)
model.optimize_graph(optimized_output_dir)
except Exception as e:
sys.exit("An error occurred during graph optimization: {}".format(str(e)))
|