def main():
# _, data_files = common.find_sample_data(description="Runs a Caffe MNIST network in Int8 mode", subfolder="mnist", find_files=["t10k-images-idx3-ubyte", "t10k-labels-idx1-ubyte", "train-images-idx3-ubyte", ModelData.DEPLOY_PATH, ModelData.MODEL_PATH])
# [test_set, test_labels, train_set, deploy_file, model_file] = data_files
#
# # Now we create a calibrator and give it the location of our calibration data.
# # We also allow it to cache calibration data for faster engine building.
# calibration_cache = "mnist_calibration.cache"
# calib = MNISTEntropyCalibrator(test_set, cache_file=calibration_cache)
#
# # Inference batch size can be different from calibration batch size.
# batch_size = 32
# with build_int8_engine(deploy_file, model_file, calib, batch_size) as engine, engine.create_execution_context() as context:
# # Batch size for inference can be different than batch size used for calibration.
# check_accuracy(context, batch_size, test_set=load_mnist_data(test_set), test_labels=load_mnist_labels(test_labels))
import requests
from PIL import Image
# IMAGE_URL = 'https://tensorflow.org/images/blogs/serving/cat.jpg'
# data = requests.get(IMAGE_URL, stream=True).content
# data = Image.open('cat.jpg')
data = batch_from_image('cat.jpg', 1)
# Convert and calibrate model
from tensorflow.python.compiler.tensorrt import trt_convert as trt
import numpy as np
import tensorflow as tf
# Load and convert a frozen graph
graph_def = tf.GraphDef()
with tf.gfile.GFile("fire4.pb", 'rb') as f:
graph_def.ParseFromString(f.read())
# dataset = tf.data.TFRecordDataset('fire-data')
# iterator = dataset.make_one_shot_iterator()
# features = iterator.get_next()
from tensorflow.python.compiler.tensorrt import trt_convert as trt
converter = trt.TrtGraphConverter(input_graph_def=graph_def,
nodes_blacklist=[
'conv_81/BiasAdd', 'conv_93/BiasAdd',
'conv_105/BiasAdd'
],
precision_mode='INT8')
graph_def = converter.convert()
def input_map_fn():
return {'input_1:0': features}
calibrated_graph_def = converter.calibrate(
fetch_names=['conv_81/BiasAdd', 'conv_93/BiasAdd', 'conv_105/BiasAdd'],
num_runs=1,
feed_dict_fn=lambda: {'input_1:0': np.array([data])})
converter.save('fire48.pb')
def from_saved_model(
tf_path: Union[Path, str],
save_path: Union[Path, str],
inputs: List[IOShape],
outputs: List[IOShape],
tf_version=2,
max_batch_size: int = 32,
max_workspace_size_bytes: int = 1 << 32,
precision_mode: str = 'FP32',
maximum_cached_engines: int = 100,
create_model_config: bool = True,
override: bool = False,
):
"""Convert TensorFlow SavedModel to TF-TRT SavedModel."""
from tensorflow.python.compiler.tensorrt import trt_convert as trt
if save_path.with_suffix('.zip').exists():
if not override: # file exist yet override flag is not set
# TODO: add logging
print('Use cached model')
return True
tf_path = Path(tf_path)
save_path = Path(save_path)
# get arch name
arch_name = parse_path(save_path)['architecture']
# TF SavedModel files should be contained in a directory
# `~/.modelci/<model-name>/tensorflow-tfs/<version>/model.savedmodel`
tf_saved_model_path = save_path / 'model.savedmodel'
assert tf_path.exists()
save_path.mkdir(parents=True, exist_ok=True)
if tf_version == 1:
converter = trt.TrtGraphConverter(
input_saved_model_dir=str(tf_path),
max_workspace_size_bytes=max_workspace_size_bytes,
precision_mode=precision_mode,
maximum_cached_engines=maximum_cached_engines
)
elif tf_version == 2:
# conversion
conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS
conversion_params = conversion_params._replace(
max_workspace_size_bytes=max_workspace_size_bytes
)
conversion_params = conversion_params._replace(precision_mode=precision_mode)
conversion_params = conversion_params._replace(
maximum_cached_engines=maximum_cached_engines
)
converter = trt.TrtGraphConverterV2(
input_saved_model_dir=str(tf_path),
conversion_params=conversion_params
)
else:
raise ValueError(f'tf_version expecting a value of `1` or `2`, but got {tf_version}')
converter.convert()
converter.save(str(tf_saved_model_path))
# zip
subprocess.call(['zip', '-r', save_path.with_suffix('.zip'), save_path])
# create model configuration
if create_model_config:
TRTConverter.generate_trt_config(
save_path.parent,
arch_name=arch_name,
platform=TensorRTPlatform.TENSORFLOW_SAVEDMODEL,
inputs=inputs,
outputs=outputs,
max_batch_size=max_batch_size
)
return True