Pytorch2onnx support return of List[Dict[str, Tensor]]

1

Hi, how to export models to onnx with output of List[Dict[str, tensor]], just like detection models in torchvision.

Following code mocking here, but does not work.
Any ideas?
Thanks in advance.

### test return with dict
import torch
import onnx
import io
import onnx.helper as helper
from torch.jit.annotations import Tuple, List, Dict, Optional
from torch import Tensor

class DummyModel(torch.nn.Module):
    def __init__(self):
        super(DummyModel, self).__init__()

    def forward(self, x: Tensor, y: Tensor, z: int)->List[Dict[str, Tensor]]:
        res = torch.jit.annotate(List[Dict[str, torch.Tensor]], [])
        xy = x + y
        xz = x + z
        yz = y + z
        res.append({'xy' : xy, 'xz' : xz, 'yz' : yz})
        return res

model = DummyModel()
input_data = (torch.arange(4).reshape(2,2) for _ in range(3))

desired = model(*input_data)
print(f'Expect:\n{desired}')

# onnx_io = "sample.onnx"
onnx_io = io.BytesIO()
torch.onnx.export(
    model,
    input_data,
    onnx_io,
    verbose=False,
    input_names=None,
    output_names=None,
)

model_onnx = onnx.load(onnx_io)

# print(helper.printable_graph(model_onnx.graph))
sess = rt.InferenceSession(out_path)
input_names = [_.name for _ in sess.get_inputs()]
print(input_names)
#forward model
input_data_npy = [_.detach().cpu().numpy() for _ in input_data]
actual = sess.run(None, {name : data for name, data in zip(input_names, input_data_npy)})
actual = actual[0]
print(f'Actual:\n{actual}')
for k, v in desired.items():
    np.testing.assert_allclose(desired[k], actual[k])
2

Hi, It seems that there is something wrong with the format of the input data.

I use input_data = [torch.arange(4).reshape(2,2) for _ in range(3)] and

torch.onnx.export(
    model,
    tuple(input_data),
    onnx_io,
    verbose=False,
    input_names=None,
    output_names=None,
)

instead, and it works fine.

3

Thanks a lot.
the output are not dict[str, tensor], but list[tensor] and the names are not right.
Do you have any suggestion with regard to this ?

Expect:
[{'xy': tensor([[0, 2],
        [4, 6]]), 'xz': tensor([[0, 2],
        [4, 6]]), 'yz': tensor([[0, 2],
        [4, 6]])}]
input_names=['0', '1', '2']
Actual:
[array([[0, 2],
       [4, 6]], dtype=int64), array([[0, 2],
       [4, 6]], dtype=int64), array([[0, 2],
       [4, 6]], dtype=int64)]
4

As discussed in github, one option is using desired_list, _ = torch.jit._flatten(desired) to set the data type of the outputs in pytorch backend to a List