test_clip_fp16.py

## 用于测试clip fp16模型，测试它为啥会溢出
# 需要对比一下，将onnx每层都标记为输出，然后转TensorRT
# 最后再将转好的trt和onnx对比，查看逐层输出结果。


#!/usr/bin/env python3
# Template auto-generated by polygraphy [v0.47.1] on 06/05/23 at 12:36:31
# code gen with onnx_trt_compare.sh
# but i edit the code to make it more readable
import tensorrt as trt
import os
import shutil
import numpy as np
from polygraphy.logger import G_LOGGER
G_LOGGER.module_severity = {'': G_LOGGER.VERBOSE}
from colored import stylize, fg
from polygraphy import constants
import onnx
from itertools import tee
from polygraphy.backend.common import SaveBytes
from polygraphy.backend.onnx import modify_outputs, onnx_from_path, ModifyOutputs
from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
from polygraphy.backend.trt import CreateConfig as CreateTrtConfig, EngineBytesFromNetwork, EngineFromBytes, ModifyNetworkOutputs, NetworkFromOnnxPath, Profile, TrtRunner
from polygraphy.common import TensorMetadata
from polygraphy.comparator import Comparator, CompareFunc, DataLoader
from polygraphy.exception import PolygraphyException
from polygraphy.backend.trt import network_from_onnx_path


# --dir info--
project_dir = os.path.dirname(os.path.abspath(__file__))
output_dir = os.path.join(project_dir, "output")
onnx_dir = os.path.join(output_dir, "onnx", "clip")
new_onnx_dir = os.path.join(output_dir, "onnx", "clip_new")
new_onnx_dir2 = os.path.join(output_dir, "onnx", "clip_new2")
for dir1 in [new_onnx_dir, new_onnx_dir2]:
    if not os.path.exists(dir1):
        os.mkdir(dir1)
    else:
        for file in os.listdir(dir1):
            os.remove(os.path.join(dir1, file))
onnx_path = os.path.join(onnx_dir,"clip.onnx")
new_onnx_path = os.path.join(new_onnx_dir,"clip.onnx")
new_onnx_path2 = os.path.join(new_onnx_dir2,"clip.onnx")
model_dir = os.path.join(output_dir, "models")
if not os.path.exists(model_dir):
    os.mkdir(model_dir)
trt_model_path = os.path.join(model_dir, "clip_fp16_mark_all.plan")
time_cache_path = os.path.join(output_dir, "clip_fp16_mark_all.cache")
use_time_cache = False


class MyLogger(trt.ILogger):
    def __init__(self):
        trt.ILogger.__init__(self)

    def log(self, severity, msg):
        if severity == trt.Logger.ERROR:
            print(stylize("[ERROR] " + msg, fg("red")))  # 红色字体
        elif severity == trt.Logger.WARNING:
            print(stylize("[WARNING] " + msg, fg("yellow")))  # 黄色字体
        elif severity == trt.Logger.INTERNAL_ERROR:
            print(stylize("[INTERNAL_ERROR] " + msg, fg("red")))  # 红色字体
        elif severity == trt.Logger.INFO:
            print(stylize("[INFO] " + msg, fg("green")))  # 绿色字体
        elif severity == trt.Logger.VERBOSE:
            print(stylize("[VERBOSE] " + msg, fg("blue")))  # 蓝色字体
        else:
            print("[UNKNOWN] " + msg)


def get_network_profiles(trt_builder):
    # ----profile1 when past_key_values is None----
    profile1 = trt_builder.create_optimization_profile()
    profile1.set_shape(
        "input_ids",
        (2, 77),
        (2, 77),
        (2, 77),
    )
    profiles = [profile1]
    return profiles


def get_network_definition(trt_network):
    def pairwise(iterable):
        a, b = tee(iterable)
        next(b, None)
        return zip(a, b)
    num_layers = trt_network.num_layers
    indices = list(range(num_layers))
    for i, i_next in pairwise(indices):
        layer = trt_network.get_layer(i)
        next_layer = trt_network.get_layer(i_next)
        layer = trt_network.get_layer(i)
        if not all([
            layer.get_output(i).is_execution_tensor
            for i in range(layer.num_outputs)
        ]):
            continue
        if layer.get_output_type(0) != trt.float32:
            continue
        if next_layer.type == trt.LayerType.SOFTMAX:
            print(layer.name, "and", next_layer.name, "force fp32.")
            layer.precision = trt.DataType.FLOAT
            next_layer.precision = trt.DataType.FLOAT
        #else:
        #    layer.precision = trt.DataType.HALF
    return trt_network


# Data Loader
dtype = np.dtype(np.int32)
data_loader = DataLoader(
    input_metadata=TensorMetadata()
    .add('input_ids', dtype=dtype, shape=(2, 77))
)
# load onnx

print("loading onnx model from", onnx_path)
onnx_model = onnx_from_path(onnx_path)
# this layer will output None in onnxrt
# bool_tensor_list = ["/transformer/encoder/layers.0/mlp/Sigmoid_output_0"]
bool_tensor_list = []
# for node in onnx_model.graph.node:
#     # print(node.name, node.op_type)
#     # this layer is a bool tensor, it will cause error when run TensorRT engine
#     if node.op_type == "Equal":
#         print("find bool", node.name)
#         bool_tensor_list.extend(node.output)
input_list = onnx_model.graph.input
input_names = [i.name for i in input_list]
output_list = onnx_model.graph.output
output_names = [o.name for o in output_list]
exclude_outputs = input_names + bool_tensor_list
# mark all layers as output for onnx model
# warning this will make the onnx model output all layers with no type and no shape
new_onnx_model = modify_outputs(
    model=onnx_model,
    outputs=constants.MARK_ALL,
    exclude_outputs=exclude_outputs,
)
new_output_list = new_onnx_model.graph.output
new_output_names = [o.name for o in new_output_list]
onnx_input_num = len(new_onnx_model.graph.input)
onnx_output_num = len(new_onnx_model.graph.output)
print("onnx input num:", onnx_input_num, "onnx output num:", onnx_output_num)
onnx.save_model(
    new_onnx_model,
    new_onnx_path,
)
# load onnx_runtime
build_onnx_rt_session = SessionFromOnnx(new_onnx_path)
# get onnx runtime output tensor info(name, type, shape)
sess = build_onnx_rt_session()
bool_tensor_list = []
sess_output = sess.get_outputs()
for node in sess_output:
    if node.type == "tensor(float)":
        dtype = onnx.TensorProto.FLOAT
    elif node.type == "tensor(float16)":
        dtype = onnx.TensorProto.FLOAT16
    elif node.type == "tensor(bool)":
        print("find bool", node.name)
        bool_tensor_list.append(node.name)
        dtype = onnx.TensorProto.BOOL
    elif node.type == "tensor(int64)":
        dtype = onnx.TensorProto.INT64
    elif node.type == "tensor(int32)":
        dtype = onnx.TensorProto.INT32
    else:
        print("unknown dtype:", node.type)
        raise ValueError
    output_tensor = onnx.helper.make_tensor_value_info(
        node.name, dtype, None
    )
    # replace the output tensor
    for i, vi in enumerate(new_onnx_model.graph.output):
        if vi.name == node.name:
            new_onnx_model.graph.output[i].CopyFrom(output_tensor)


# change onnx -inf to -1e4
for node in new_onnx_model.graph.node:
    # if node.name == "/text_model/ConstantOfShape_1":
    if node.op_type == "ConstantOfShape":
        print(node)
        attr = node.attribute[0]
        print(attr)
        if attr.name == "value" and attr.t.data_type == onnx.TensorProto.FLOAT:
            np_array = np.frombuffer(attr.t.raw_data, dtype=np.float32).copy()
            print("raw array", np_array)
            np_array[np_array == -np.inf] = -100000  # 将所有负无穷的值改为-100000
            attr.t.raw_data = np_array.tobytes() 
            print("new array", np_array)
        print(attr)
onnx.save_model(
    new_onnx_model,
    new_onnx_path2,
    # save_as_external_data=True,
    # all_tensors_to_one_file=False
)

print("===========onnx model loaded=========================")
# build trt engine
print("===========building trt engine=========================")
# prepare trt builder
builder = trt.Builder(MyLogger())
builder.max_threads = os.cpu_count() // 2
config = builder.create_builder_config()
# use fp16
config.flags = 1 << int(trt.BuilderFlag.FP16)
# disable TF32
config.flags = config.flags & ~(1 << int(trt.BuilderFlag.TF32))
# config.flags = config.flags | (1 << int(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS))
print("use fp16? ", config.get_flag(trt.BuilderFlag.FP16))
print("use obey? ", config.get_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS))
# use obey precision constraints
# read time cache
if use_time_cache:
    if os.path.exists(time_cache_path):
        time_cache = open(time_cache_path, "rb").read()
        if time_cache is None:
            time_cache = b""
            print(stylize("read time cache failed", fg("red")))
        else:
            print(stylize(f"read time cache from {time_cache_path}", fg("green")))
    else:
        time_cache = b""
        print(stylize("time cache will init with empty.", fg("green")))

    # set time cache
    cache = config.create_timing_cache(time_cache)
    config.set_timing_cache(cache, False)
profile_list = get_network_profiles(builder)
for profile in profile_list:
    config.add_optimization_profile(profile)

_b, network, _p = NetworkFromOnnxPath(new_onnx_path2)()
# network = get_network_definition(network)

# network = make_logsoftmax_fp32(network)

# network = network_from_onnx_path(onnx_path)
# set_network_outputs = ModifyNetworkOutputs(
#     network=network,
#     outputs=constants.MARK_ALL,
#     exclude_outputs=bool_tensor_list
# )
# wo_b, network, _p = set_network_outputs()
network_input_number = network.num_inputs
network_output_number = network.num_outputs
print("TensorRT input num:", network_input_number, "TensorRT output num:", network_output_number)
serialized_engine = builder.build_serialized_network(network, config)
if serialized_engine is not None:
    with open(trt_model_path, "wb") as f:
        f.write(serialized_engine)
    # save_engine(trt_engine, tensorrt_engine_path)
    print("==tensorRT engine compile done==")

# save time cache
if use_time_cache and not os.path.exists(time_cache_path):
    time_cache = config.get_timing_cache()
    if time_cache is not None:
        time_cache_data = time_cache.serialize()
        open(time_cache_path, "wb").write(time_cache_data)
        print(
            stylize(
                "save time cache to {}".format(time_cache_path),
                fg("green")
            )
        )

# profiles = [
#     Profile().add('input_ids:[1,512] position_ids:[1,2,512] attention_mask:[1,1,512,512] past_key_values.0.decorder.key:[0,1,32,128] past_key_values.0.decorder.value', min=[0, 1, 32, 128], opt=[0, 1, 32, 128], max=[0, 1, 32, 128])
# ]
# create_trt_config = CreateTrtConfig(memory_pool_limits=2 * (1024 ** 3), profiles=profile_list)
# build_engine = EngineBytesFromNetwork(set_network_outputs, config=create_trt_config)
# save_engine_bytes = SaveBytes(build_engine, path=trt_model_path)
save_engine_bytes = SaveBytes(serialized_engine, path=trt_model_path)
deserialize_engine = EngineFromBytes(save_engine_bytes)
print("===========trt engine build OK=========================")

# Runners
runners = [
    OnnxrtRunner(build_onnx_rt_session),
    TrtRunner(deserialize_engine),
]

# Runner Execution
results = Comparator.run(runners, data_loader=data_loader)

success = True
# Accuracy Comparison
compare_func = CompareFunc.simple(rtol={'': 5e-2}, atol={'': 5e-2})
success &= bool(Comparator.compare_accuracy(results, compare_func=compare_func))

# Report Results
if not success:
    raise PolygraphyException('FAILED')