pytorch / pytorch/vision

ONNX export of MaskRCNN: dynamic axes seem broken (for batch size > 1)

Open
#2,311 6 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug module: onnx topic: object detection
Dominant language
Python
Stars
17.9k
Forks
7.3k
Avg merge
1d 15h
Merged PRs (30d)
13

Description

🐛 Bug

At export to ONNX, dynamic axes were set and the inputs and outputs named properly. However, the output of inferred images is incorrect and wrongly named. Depending on different batch sizes used at export and inference, the behaviour varies as follows:
Supposing that batch size at export time is n, and batch size at inference time is m:

  1. if n==m:
    Output has length of n*4, so ex. if n=m=3, output has length of 12. In the onnx runtime session, it looks like in the following screenshot:
    image
    So only output of the first image in batch is correctly named.
  2. if n<m:
    Similar behaviour as in 1., but output for only n first images in batch is returned.
  3. if n>m:
    A "SplitToSequence_XXXX" error is returned, ex. that one:
onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running SplitToSequence node. Name:'SplitToSequence_4001' Status Message: split_size_sum (57) != split_dim_size (40)

This exception is similar to behaviour listed in #2309 and seems connected.

To Reproduce

Steps to reproduce the behavior:
1.
Load and export a pretrained MaskRCNN model using input_tensor of shape (n, 3, 1024, 1024), for example (4,3,1024,1024):

model = torchvision.models.detection.maskrcnn_resnet50_fpn(
            pretrained=False,
            min_size=1024, max_size=1024,
            pretrained_backbone=False,
            num_classes=num_classnames + 1,  # + background class
            image_mean=image_mean,
            image_std=image_std,
)
torch.onnx.export(
        model,
        input_tensor.float(),
        onnx_model_filepath,
        export_params=True,
        opset_version=12,
        do_constant_folding=False,
        input_names=["images_tensors"],
        output_names=["boxes", "labels", "scores", "masks"],
        dynamic_axes={"images_tensors": [0, 1, 2, 3], "boxes": [0, 1], "labels": [0],
                      "scores": [0], "masks": [0, 1, 2, 3]},
)
  1. Load and infer ONNX model on input_tensor of shape (m,3,1024,1024), where m corresponds to the value in the description above, and different m values (bigger, smaller or equal to n) will result in different behaviours.
input_array = input_tensor.cpu().numpy()
ort_session = onnxruntime.InferenceSession(onnx_model_filepath)
ort_inputs = {"images_tensors": input_array}
ort_outs = ort_session.run(None, ort_inputs)
outputs = ort_session.get_outputs()

These outputs are presented in the screenshot above.

Expected behavior

With dynamic_axes set properly, I expect:

  1. Output length dependent on batch size of the inferred tensor, not the one used for export.
  2. Output of shape similar to the torch model's output, which is a list (len == batch size) of dictionaries of boxes, labels, scores and masks. Also, all outputs correctly named, not like currently in the above screenshot.
  3. No exceptions if inferred tensor is of smaller batch size than the one used for export.

Environment

PyTorch version: 1.6.0.dev20200526+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 20.04 LTS
GCC version: (Ubuntu 9.3.0-10ubuntu2) 9.3.0
CMake version: version 3.16.3

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti

Nvidia driver version: 440.64
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.4

Versions of relevant libraries:
[pip3] numpy==1.18.4
[pip3] torch==1.6.0.dev20200526+cu101
[pip3] torchvision==0.7.0.dev20200526+cu101
[conda] Could not collect

Also:

ONNX_runtime and ONNX_runtime_gpu==1.3.0
ONNX==1.7.0

Additional context

This seems connected to #2309 and #2251

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the torch.onnx.export and ONNX Runtime workflow described in the issue with differing export and inference batch sizes. Read the related issues #2309 and #2251 for context. Done means dynamic inference returns correctly named per-image MaskRCNN outputs for varying batch sizes without SplitToSequence errors.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
computer-vision, machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.