I had a related problem with Faster R-CNN and onnxruntime. When I export the model using a dummy input which doesn’t contain an object, at inference time the model works only with similar kind of images. Meaning, if there is no detection it is fine, but If the model detects an object, it crashes with the message The input tensor cannot be reshaped to the requested shape. Same thing happens vice versa. If I export the model using a real image, than onnxruntime crashes when there is no detection at inference time (empty tensor). So, at some point during the export the reshape function is hardcoded for a certain tensor shape and can’t handle both detection and no-detection cases. My solution was that I concatenated the predictions with zeros, so that even though the predictions is an empty tensor, it would now contain bunch of zeros, and reshape would take place safely. However, I had to use a dummy input (which leads to no detection) during export for this scheme to work. I added the following code at this line: vision/roi_heads.py at 3e27eb2104c01a7a55f5b4e82d19d9b6612715b5 · pytorch/vision · GitHub 5, right before pred_boxes_list = ...
# Added to fix "reshape mismatch error" when there is no detection
pred_boxes = torch.cat((pred_boxes.view(-1, num_classes, 4), torch.zeros((1, num_classes, 4), device=device)), dim=0)
pred_scores = torch.cat((pred_scores, torch.zeros((1, num_classes), device=device)), dim=0)
boxes_per_image[0] += 1