Description
Dynamic Input. IShuffleLayer applied to shape tensor must have 0 or 1 reshape dimensions
Environment
TensorRT Version: 8.5.2-1+cuda11.8
GPU Type: Tesla K80
Nvidia Driver Version: 470.161.03
CUDA Version: 11.7
CUDNN Version:
Operating System + Version: Ubuntu 20.04.5 LTS x86_64
Python Version (if applicable): 3.9
TensorFlow Version (if applicable):
PyTorch Version (if applicable): 1.10.2
Baremetal or Container (if container which image + tag):
Relevant Files
Please attach or include links to any models, data, files, or scripts necessary to reproduce your issue. (Github repo, Google Drive, Dropbox, etc.)
Steps To Reproduce
Imports
import torch
import torchvision
import torch.nn.functional as F
Set the dummy input to input size
dummy_input = torch.randn(1, 3, 224, 224)
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
Providing input and output names sets the display names for values within the model’s graph.
input_names = [ “input” ]
output_names = [ “output” ]
class TraceWrapper(torch.nn.Module):
def init(self, model):
super().init()
self.model = model
def forward(self, x):
out = self.model(x)
boxes = out[0][“boxes”].reshape(-1)
labels = out[0]["labels"].reshape(-1)
scores = out[0][“scores”].reshape(-1)
return labels
# return torch.rand(1,2)
model = TraceWrapper(model)
prediction = model(torch.rand(1,3,224,224))
print(prediction)
The call to torch.onnx.export runs the model once to trace its execution and then exports the traced model to the specified file.
torch.onnx.export(model, dummy_input, “model.onnx”, verbose=True,
input_names=input_names, output_names=output_names, opset_version=13)
Import tensorrt module to access the Python API
import tensorrt as trt
Create a logger
logger = trt.Logger(trt.Logger.WARNING)
Create a builder
builder = trt.Builder(logger)
Create a build configuration specifying how TensorRT should optimize the model
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20) # 1 MiB
Create a network definition:
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
Create an ONNX parser to populate the network as follows:
parser = trt.OnnxParser(network, logger)
Read the model file and process any errors:
success = parser.parse_from_file(“folded.onnx”)
for idx in range(parser.num_errors):
print(parser.get_error(idx))
if not success:
pass # Error handling code here
Engine can be built and serialized with
serialized_engine = builder.build_serialized_network(network, config)
Save the engine to a file for future use
with open(“sample.engine2”, “wb”) as f:
f.write(serialized_engine)