-
Notifications
You must be signed in to change notification settings - Fork 375
Description
Is your feature request related to a problem? Please describe.
I am converting an ONNX model into TensorRT mixed-precision. because of the bad generation result, we had to tell the converion script where of the model should be kept fp32 precision.
We got every layer output of the ONNX model, and find out all outputs that with large amptitude. So now we had a list of dangerous ONNX output name.
Then we set the layers in the script:
for in range(network.num_layers):
layer = network.get_layer(i)
for j in range(layer.num_outputs):
output_name = layer.get_output(j).name
if output_name in dangeous_outputs:
layer.precison = trt.float32
layer.set_output_type(j, trt.float32)
but this time, we found that some types of layers are not resetable. for example, the layere with type trt.LayerType.SLICE
So then, we had to found out which layerType are forbidden from being set as fp32.
That is disgusting and costly. If that layer can not be set, just ignore please, do not broken my converion process.
Describe the solution you'd like
just provice simple and lazy API for model layer precision lock. get ride of the details. eg:
def trt.lock_precision_by_layername() # just for the whole layer name precison locking. like /transformer1/ff0/layernorm
def trt.lock_precison_by_outputname() # just for the layer with the specified outputname
def trt.lock_precison_by_layername_pattern() # just for the layer whose name match the specified pattern, like '/attention/'
Describe alternatives you've considered
Additional context