wehub-resource-sync
262 lines
11 KiB
Python
262 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# 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.
|
|
#
|
|
import argparse
|
|
import onnx
|
|
import logging
|
|
import os
|
|
import numpy as np
|
|
import onnx_graphsurgeon as gs
|
|
from onnx import helper, TensorProto, numpy_helper, version_converter
|
|
|
|
'''
|
|
This script is converting TE ONNX models (cast + CustomOp Q) and (CustomOp DQ + cast) pairs to Opset19 ONNX Q/DQ
|
|
usage:
|
|
python3 convert_te_onnx_to_trt_onnx.py --onnx_model_path <folder/file>
|
|
|
|
This script requires onnx 1.14 and above
|
|
'''
|
|
|
|
def find_node_by_tensor(graph, search_tensor, is_node_input, search_node_type=None):
|
|
for idx, node in enumerate(graph.node):
|
|
search_container = node.output
|
|
if is_node_input:
|
|
search_container = node.input
|
|
for node_tensor in search_container:
|
|
if search_node_type and node.op_type != search_node_type:
|
|
continue
|
|
if node_tensor == search_tensor:
|
|
return node, idx
|
|
|
|
return None, None
|
|
|
|
def redirect_quantize_input(graph, q_node):
|
|
assert(q_node.op_type == 'QuantizeLinear')
|
|
q_input = q_node.input[0]
|
|
cast_node, cast_node_idx = find_node_by_tensor(graph, q_input, False, 'Cast')
|
|
if cast_node:
|
|
q_node.input[0] = cast_node.input[0]
|
|
return [cast_node_idx]
|
|
return []
|
|
|
|
def redirect_dequantize_output(graph, dq_node):
|
|
assert(dq_node.op_type == 'DequantizeLinear')
|
|
dq_output = dq_node.output[0]
|
|
cast_node, cast_node_idx = find_node_by_tensor(graph, dq_output, True, 'Cast')
|
|
if cast_node:
|
|
dq_node.output[0] = cast_node.output[0]
|
|
return [cast_node_idx]
|
|
return []
|
|
|
|
def get_attr_numpy_tensor(attr):
|
|
assert(attr.type == onnx.AttributeProto.TENSOR)
|
|
return numpy_helper.to_array(attr.t)
|
|
|
|
def get_attr(node, search_attr_name):
|
|
for idx, attr in enumerate(node.attribute):
|
|
if attr.name == search_attr_name:
|
|
return attr, idx
|
|
|
|
return None, None
|
|
|
|
def cast_scale(graph, qdq_node, cast_to):
|
|
assert(cast_to in ['fp32', 'fp16'])
|
|
assert(qdq_node.op_type in ['QuantizeLinear', 'DequantizeLinear'])
|
|
constant_node_idx = None
|
|
scale_tensor = qdq_node.input[1]
|
|
constant_node, constant_node_idx = find_node_by_tensor(graph, scale_tensor, False, 'Constant')
|
|
scale_cast_to_dtype = None
|
|
onnx_cast_to_dtype = None
|
|
if cast_to == 'fp16':
|
|
scale_cast_to_dtype = np.dtype(np.float32)
|
|
onnx_cast_to_dtype = onnx.TensorProto.FLOAT16
|
|
elif cast_to == 'fp32':
|
|
scale_cast_to_dtype = np.dtype(np.float32)
|
|
onnx_cast_to_dtype = onnx.TensorProto.FLOAT
|
|
|
|
if constant_node:
|
|
scale_attr, _ = get_attr(constant_node, 'value')
|
|
assert(scale_attr)
|
|
numpy_scale = get_attr_numpy_tensor(scale_attr)
|
|
logging.info(type(numpy_scale.dtype))
|
|
logging.info(type(scale_cast_to_dtype))
|
|
if numpy_scale.dtype != scale_cast_to_dtype:
|
|
logging.debug(f'Change {qdq_node.name} scale from {numpy_scale.dtype} to {scale_cast_to_dtype}')
|
|
numpy_scale = numpy_scale.astype(scale_cast_to_dtype)
|
|
tensor_name = constant_node.name + '_casted'
|
|
create_constant_tensor(graph, tensor_name, onnx_cast_to_dtype, numpy_scale)
|
|
qdq_node.input[1] = tensor_name
|
|
else:
|
|
logging.warning(f'No constant node connected to {qdq_node} as scale')
|
|
|
|
if constant_node_idx:
|
|
return [constant_node_idx]
|
|
return []
|
|
|
|
def create_constant_tensor(graph, name, dtype, np_tensor):
|
|
tensor_value_info = helper.make_tensor_value_info(name, dtype, np_tensor.shape)
|
|
graph.input.append(tensor_value_info)
|
|
helper.make_tensor(name, data_type=dtype, dims=(), vals=[0])
|
|
|
|
tensor_initializer = helper.make_tensor(name, dtype, np_tensor.shape, np_tensor.flatten().tolist())
|
|
graph.initializer.append(tensor_initializer)
|
|
|
|
'''
|
|
Convert custom operators to opset19
|
|
'''
|
|
def custom_op_to_opset19(graph, node, use_int32_quantization, remove_cast_before_q, remove_cast_after_dq, change_qdq_scale_precision):
|
|
assert(node.op_type in ['TRT_FP8QuantizeLinear', 'TRT_FP8DequantizeLinear'])
|
|
is_dq = node.op_type == 'TRT_FP8DequantizeLinear'
|
|
logging.debug(f'Convert {node.name} to Opset19')
|
|
orig_node_name = node.name
|
|
new_node_name = orig_node_name + '_converted'
|
|
|
|
quant_to = TensorProto.FLOAT8E4M3FN
|
|
if use_int32_quantization:
|
|
quant_to = TensorProto.INT32
|
|
|
|
#add zero point to the node
|
|
tensor_name = new_node_name + '_zero_point'
|
|
create_constant_tensor(graph, tensor_name, quant_to, np.array([0]))
|
|
node.input.append(tensor_name)
|
|
|
|
node.domain = ""
|
|
node.op_type = "QuantizeLinear"
|
|
|
|
node_idxs_to_delete = []
|
|
if is_dq:
|
|
node.op_type = "DequantizeLinear"
|
|
if remove_cast_after_dq:
|
|
node_idxs_to_delete += redirect_dequantize_output(graph, node)
|
|
if change_qdq_scale_precision:
|
|
node_idxs_to_delete += cast_scale(graph, node, change_qdq_scale_precision)
|
|
else:
|
|
if remove_cast_before_q:
|
|
node_idxs_to_delete += redirect_quantize_input(graph, node)
|
|
if change_qdq_scale_precision:
|
|
node_idxs_to_delete += cast_scale(graph, node, change_qdq_scale_precision)
|
|
|
|
node.name = new_node_name
|
|
logging.debug(f'Convert Done\n')
|
|
return node_idxs_to_delete
|
|
|
|
def check_model(graph):
|
|
converted_qdq_ops = ['TRT_FP8QuantizeLinear', 'TRT_FP8DequantizeLinear']
|
|
passed_check = True
|
|
for node in graph.node:
|
|
if node.op_type in converted_qdq_ops:
|
|
logging.error(f'Node \"{node.name}\" of type {node.op_type} should have been removed')
|
|
passed_check = False
|
|
return passed_check
|
|
|
|
def update_quantize_node_type(model):
|
|
graph = gs.import_onnx(model)
|
|
for node in graph.nodes:
|
|
if node.op == "TRT_FP8QuantizeLinear":
|
|
for out in node.outputs:
|
|
out.dtype = TensorProto.FLOAT8E4M3FN
|
|
return gs.export_onnx(graph)
|
|
|
|
'''
|
|
Converts onnx files from TE to TRT
|
|
'''
|
|
def replace_customop_qdq_with_onnx_qdq(te_onnx_files, results_path, create_netron_compatible_model, remove_cast_before_q, remove_cast_after_dq, change_qdq_scale_precision):
|
|
# store mappings from original ONNX name to new ONNX name.
|
|
file_mappings = {}
|
|
for te_onnx_file in te_onnx_files:
|
|
logging.debug('Loading model')
|
|
model = onnx.load(te_onnx_file, load_external_data=False)
|
|
# update QuantizeLinear output dtype
|
|
model = update_quantize_node_type(model)
|
|
# change model opset to 19
|
|
model.opset_import[0].version = 19
|
|
graph = model.graph
|
|
logging.debug('Loading model finished')
|
|
converted_qdq_ops = ['TRT_FP8QuantizeLinear', 'TRT_FP8DequantizeLinear']
|
|
|
|
try:
|
|
node_idxs_to_delete = []
|
|
converted = False
|
|
for node in graph.node:
|
|
if node.op_type in converted_qdq_ops:
|
|
converted = True
|
|
node_idxs_to_delete += custom_op_to_opset19(graph, node, create_netron_compatible_model, remove_cast_before_q, remove_cast_after_dq, change_qdq_scale_precision)
|
|
|
|
if converted:
|
|
assert(check_model(graph))
|
|
node_idxs_to_delete = reversed(sorted(node_idxs_to_delete))
|
|
for node_idx in node_idxs_to_delete:
|
|
del(graph.node[node_idx])
|
|
suffix = '.opset19'
|
|
if create_netron_compatible_model:
|
|
suffix += '.netron'
|
|
suffix += '.onnx'
|
|
new_model_filename = os.path.join(results_path, os.path.splitext(os.path.split(te_onnx_file)[1])[0] + suffix)
|
|
onnx.save_model(model, new_model_filename)
|
|
logging.info(f'The converted model is saved at {new_model_filename}!')
|
|
file_mappings[te_onnx_file] = new_model_filename
|
|
else:
|
|
logging.info(f'No conversion was done with {te_onnx_file}!')
|
|
file_mappings[te_onnx_file] = te_onnx_file
|
|
except Exception as ex:
|
|
logging.error(f'Failed: {ex}')
|
|
file_mappings[te_onnx_file] = None
|
|
return file_mappings
|
|
|
|
if __name__ == "__main__":
|
|
logging.getLogger().setLevel(logging.INFO)
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--onnx_model_path', required=True, help="Path of model or a folder of models. When using a folder, this script will convert all \'.onnx\' files")
|
|
parser.add_argument('--results_path', required=False, help="Path for generated models, when not set, the generated model(s) will be next ot the origianl model(s)")
|
|
parser.add_argument('--create_netron_compatible_model', action='store_true', required=False, help="When set, the script will use int32 quantization. "
|
|
"This enables the user to view the graph with Netron, until it adds support for opset19. The generated model isn't TRT compatible.")
|
|
parser.add_argument('--remove_casts', required=False, help="Controls whether to remove casts around q/dq nodes. "
|
|
"For example, when set to \'dq\', remove casts only after dq. Default is \'keep_all\'", choices=['q', 'dq', 'qdq', 'keep_all'], default='keep_all')
|
|
parser.add_argument('--change_qdq_scale_precision', required=False, help="When set controls q/dq nodes scales data type.", choices=['fp32', 'fp16'])
|
|
args = parser.parse_args()
|
|
|
|
results_path = args.results_path
|
|
if results_path and os.path.isdir(results_path) == False:
|
|
logging.error(f'\'--results_path\' set to \'{results_path}\', but the folder doesn\'t exist, exiting')
|
|
exit(-1)
|
|
|
|
if results_path is None:
|
|
results_path = args.onnx_model_path
|
|
if os.path.isfile(results_path):
|
|
results_path = os.path.split(results_path)[0]
|
|
|
|
remove_cast_after_dq = False
|
|
remove_cast_before_q = False
|
|
if args.remove_casts == 'q':
|
|
remove_cast_before_q = True
|
|
elif args.remove_casts == 'dq':
|
|
remove_cast_after_dq = True
|
|
elif args.remove_casts == 'qdq':
|
|
remove_cast_after_dq = True
|
|
remove_cast_before_q = True
|
|
|
|
onnx_files = []
|
|
if os.path.isdir(args.onnx_model_path):
|
|
logging.info(f"Got folder: {args.onnx_model_path}")
|
|
onnx_files = [os.path.join(args.onnx_model_path, filename) for filename in os.listdir(args.onnx_model_path) if filename.endswith('.onnx')==True and filename.endswith('.opset19.onnx')==False]
|
|
|
|
else:
|
|
logging.info(f"Got file: {args.onnx_model_path}")
|
|
onnx_files = [args.onnx_model_path]
|
|
|
|
replace_customop_qdq_with_onnx_qdq(onnx_files, results_path, args.create_netron_compatible_model, remove_cast_before_q, remove_cast_after_dq, args.change_qdq_scale_precision)
|