Description
I have a channel last TF model, and I convert it to onnx → trt. When invoking trtexec, even if I set --inputIOFormats=fp32:hwc, the input is still handled as channel first, and a pair of transposes (from channel last to channel first, then from channel first to channel last) are added. I wonder how I can get rid of these transposes to get better performance?
Environment
TensorRT Version: 8.5.1
GPU Type: RTX4000
Nvidia Driver Version: 525
CUDA Version: 11.8
CUDNN Version: 8.6
Operating System + Version: Ubuntu 20.04
Python Version (if applicable): 3.8
TensorFlow Version (if applicable): 2.12
PyTorch Version (if applicable):
Baremetal or Container (if container which image + tag):
Steps To Reproduce
- Run this Python file
import os
import tensorflow as tf
import numpy as np
SAVED_MODEL_DIR = "/tmp/resnet50"
ONNX_MODEL_PATH = SAVED_MODEL_DIR + ".onnx"
class ResNet50(tf.Module):
def __init__(self):
super().__init__()
self.model = tf.keras.applications.resnet50.ResNet50(
weights="imagenet",
include_top=True
)
@tf.function
def forward(self, inputs):
return self.model(inputs, training=False)
resnet = ResNet50()
input_batch = np.float32(np.random.rand(1, 224, 224, 3))
print("tf result", resnet.forward(input_batch)[0, 0])
# Save saved model.
tensor_specs = [tf.TensorSpec((1, 224, 224, 3), tf.float32)]
call_signature = resnet.forward.get_concrete_function(*tensor_specs)
os.makedirs(SAVED_MODEL_DIR, exist_ok=True)
print(f"Saving {SAVED_MODEL_DIR} with call signature: {call_signature}")
tf.saved_model.save(resnet, SAVED_MODEL_DIR,
signatures={"serving_default": call_signature})
# convert to onnx
assert os.system(
f"python -m tf2onnx.convert --saved-model {SAVED_MODEL_DIR} --output {ONNX_MODEL_PATH}") == 0
# convert to trt
assert os.system(f"trtexec --onnx={ONNX_MODEL_PATH} --verbose --inputIOFormats=fp32:hwc") == 0
- From the log, we can see transpose is added
StatefulPartitionedCall/resnet50/conv1_conv/Conv2D__6 [Transpose] inputs: [inputs -> (1, 224, 224, 3)[FLOAT]],
- Remove
--inputIOFormats=fp32:hwcand rerun, you can get an exactly the same engine, which means it doesn’t take effect.
