[Torch] Add support for aten::dot, aten::mv and prim::DictConstruct

[pseudomo]

I am trying to optimize this model https://github.com/yucornetto/MGMatting
...
model.load_state_dict(...)
model = model.eval()

model = torch.jit.trace(model, (image, mask), strict=False).eval()
mod, params = relay.frontend.from_pytorch(model, [('image', image.shape), ('guidance', mask.shape)])
...

But
relay.frontend.from_pytorch fails with the following error:
The following operators are not implemented: ['aten::dot', 'prim::DictConstruct', 'aten::mv']

I could get rid of prim::DictConstruct operation but aten::dot and aten:mv are needed

Environment

Operating System: Windows 10
Torch version: 1.10.1+cpu
TVM version: git commit 115919b (18 Dec 2021)

[masahi]

Can you share your complete test script?

[pseudomo]

@masahi

Can you share your complete test script?

Assuming that TVM and all requirements for MGMatting are installed

git clone https://github.com/yucornetto/MGMatting.git
cd MGMatting/code-base
mkdir pretrain
gdown --id 1RO2JOvPz1YzeHoMCTojNwkceYZamFZ7Q -O pretrain/latest_model.pth

Create to_tvm.py with the following content:

import os
import torch
import utils
from utils import CONFIG
import networks
from tvm import relay

current_unit_dir = os.path.dirname(os.path.realpath(file))
model = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder)
model.cpu()
checkpoint = torch.load(os.path.join(current_unit_dir,'pretrain/latest_model.pth'), map_location=torch.device('cpu'))
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)

model = model.eval()

image = torch.rand(1,3,1024,1024)
mask = torch.rand(1,1,1024,1024)
model = torch.jit.trace(model, (image, mask), strict=False)
mod, params = relay.frontend.from_pytorch(model, [('image', image.shape), ('guidance', mask.shape)])

This code will raise the exception:
Traceback (most recent call last):
File ".\to_tvm.py", line 20, in
mod, params = relay.frontend.from_pytorch(model, [('image', image.shape), ('guidance', mask.shape)])
File "<...>\anaconda3\envs\tvm-build\lib\site-packages\tvm-0.9.dev196+g115919b12-py3.7-win-amd64.egg\tvm\relay\frontend\pytorch.py", line 3886, in from_pytorch
converter.report_missing_conversion(op_names)
File "<...>\anaconda3\envs\tvm-build\lib\site-packages\tvm-0.9.dev196+g115919b12-py3.7-win-amd64.egg\tvm\relay\frontend\pytorch.py", line 3119, in report_missing_conversion
raise NotImplementedError(msg)
NotImplementedError: The following operators are not implemented: ['aten::mv', 'prim::DictConstruct', 'aten::dot']

[Leeviber]

Hi, I met same issue with you, how did you get rid of prim::DictConstruct operation? I only stuck with this op.

[zyc-bit]

Hi, I met same issue with you, how did you get rid of prim::DictConstruct operation? I only stuck with this op.

maybe change the code in your model, such as use tuple instand of dict......