KeyError: 'unbind'

[schyun9212]

🐛 Bug

To Reproduce

Steps to reproduce the behavior:

# %%
import torch

from PIL import Image
from maskrcnn_benchmark.config import cfg
from predictor import COCODemo
from maskrcnn_benchmark.structures.image_list import ImageList

from demo.utils import load_image, imshow, masking_image
from demo.transform import transform_image
import os

OVERWRITE_MODEL = True
TEST_IMAGE_PATH = "./sample.jpg"
# MODEL_DEVICE = "cuda"
MODEL_DEVICE = "cpu"
MODEL_PATH = f"./maskrcnn_{MODEL_DEVICE}.onnx"
ONNX_OPSET_VERSION = 10

config_file = "../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml"
cfg.merge_from_file(config_file)
cfg.merge_from_list(["MODEL.DEVICE", MODEL_DEVICE])
cfg.freeze()

coco_demo = COCODemo(
    cfg,
    confidence_threshold=0.7,
    min_image_size=800,
)

class MaskRCNNModel(torch.nn.Module):
    def __init__(self):
        super(MaskRCNNModel, self).__init__()
        for param in coco_demo.model.parameters():
            param.requires_grad = False

    def forward(self, image):
        image_list = ImageList(image.unsqueeze(0), [(int(image.size(-2)), int(image.size(-1)))])

        result, = coco_demo.model(image_list)

        result = (result.bbox,
                result.get_field("labels"),
                result.get_field("mask"),
                result.get_field("scores"))
        return result

original_image = load_image(TEST_IMAGE_PATH)
image, t_width, t_height = transform_image(cfg, original_image)

height, width = original_image.shape[:-1]

# requires_grad must be false for tracing
if OVERWRITE_MODEL or not os.path.exists(MODEL_PATH):
    model = MaskRCNNModel()
    model.eval()
    torch.onnx.export(model, (image, ), MODEL_PATH,
                        do_constant_folding=True, opset_version=ONNX_OPSET_VERSION)

Expected behavior

KeyError                                  Traceback (most recent call last)
~/Workspace/maskrcnn/maskrcnn-benchmark/demo/export_to_onnx.py in 
     55     model.eval()
     56     torch.onnx.export(model, (image, ), MODEL_PATH,
---> 57                         do_constant_folding=True, opset_version=ONNX_OPSET_VERSION)

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/__init__.py in export(model, args, f, export_params, verbose, training, input_names, output_names, aten, export_raw_ir, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
    141                         operator_export_type, opset_version, _retain_param_name,
    142                         do_constant_folding, example_outputs,
--> 143                         strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
    144 
    145 

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/utils.py in export(model, args, f, export_params, verbose, training, input_names, output_names, aten, export_raw_ir, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
     64             _retain_param_name=_retain_param_name, do_constant_folding=do_constant_folding,
     65             example_outputs=example_outputs, strip_doc_string=strip_doc_string,
---> 66             dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
     67 
     68 

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, propagate, opset_version, _retain_param_name, do_constant_folding, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size)
    380                                                         example_outputs, propagate,
    381                                                         _retain_param_name, do_constant_folding,
--> 382                                                         fixed_batch_size=fixed_batch_size)
    383 
    384         # TODO: Don't allocate a in-memory string for the protobuf

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, training, input_names, output_names, operator_export_type, example_outputs, propagate, _retain_param_name, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size)
    260     graph = _optimize_graph(graph, operator_export_type,
    261                             _disable_torch_constant_prop=_disable_torch_constant_prop,
--> 262                             fixed_batch_size=fixed_batch_size)
    263 
    264     if isinstance(model, torch.jit.ScriptModule) or isinstance(model, torch.jit.Function):

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/utils.py in _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size)
    130         torch._C._jit_pass_erase_number_types(graph)
    131 
--> 132         graph = torch._C._jit_pass_onnx(graph, operator_export_type)
    133         torch._C._jit_pass_lint(graph)
    134 

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/__init__.py in _run_symbolic_function(*args, **kwargs)
    172 def _run_symbolic_function(*args, **kwargs):
    173     from torch.onnx import utils
--> 174     return utils._run_symbolic_function(*args, **kwargs)
    175 
    176 

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/utils.py in _run_symbolic_function(g, n, inputs, env, operator_export_type)
    616                                   "torch.onnx.symbolic_opset{}.{} does not exist"
    617                                   .format(op_name, opset_version, op_name))
--> 618                 op_fn = sym_registry.get_registered_op(op_name, '', opset_version)
    619                 return op_fn(g, *inputs, **attrs)
    620 

~/.pyenv/versions/maskrcnn-tracing-latest/lib/python3.7/site-packages/torch/onnx/symbolic_registry.py in get_registered_op(opname, domain, version)
     89         warnings.warn("ONNX export failed. The ONNX domain and/or version are None.")
     90     global _registry
---> 91     return _registry[(domain, version)][opname]

KeyError: 'unbind'

Environment

PyTorch version: 1.3.1 Is debug build: No CUDA used to build PyTorch: 10.1.243

OS: Ubuntu 18.04.3 LTS GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0 CMake version: version 3.10.2

Python version: 3.7 Is CUDA available: Yes CUDA runtime version: 10.1.243 GPU models and configuration: GPU 0: GeForce RTX 2080 Ti Nvidia driver version: 440.44 cuDNN version: Probably one of the following: /usr/local/cuda-10.0/targets/x86_64-linux/lib/libcudnn.so.7 /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5

Versions of relevant libraries: [pip3] numpy==1.18.1 [pip3] torch==1.3.1 [pip3] torchvision==0.4.2 [conda] Could not collect

[schyun9212]

There is no concept of tensor list in ONNX. Without this concept, it is very hard to export operators that consume or produce tensor list, especially when the length of the tensor list is not known at export time.

x = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])

# This is not exportable
class Model(torch.nn.Module):
    def forward(self, x):
        return x.unbind(0)

# This is exportable.
# Note that in this example we know the split operator will always produce exactly three outputs,
# Thus we can export to ONNX without using tensor list.
class AnotherModel(torch.nn.Module):
    def forward(self, x):
        return [torch.squeeze(out, 0) for out in torch.split(x, [1,1,1], dim=0)]

https://pytorch.org/docs/stable/onnx.html