[Pytroch] prim::PythonOp' is not implemented

I ran into this problem

Traceback (most recent call last):

  File "src/test.py", line 38, in <module>
    main()

  File "src/test.py", line 35, in main
    mod, params = relay.frontend.from_pytorch(scripted_model, {'input.1': img_shape, 'input.2':boxes_shape})

  File "/data00/laiqingquan/tvm/python/tvm/relay/frontend/pytorch.py", line 1025, in from_pytorch
    _report_missing_conversion(op_names)

  File "/data00/laiqingquan/tvm/python/tvm/relay/frontend/pytorch.py", line 790, in _report_missing_conversion
    raise NotImplementedError(msg)

NotImplementedError: The following operators are not implemented: ['prim::PythonOp']

When my code was similar to import os import sys

import tvm
import torch
import torchvision
import numpy as np
from torch import nn
from tvm import relay

class Wrap(nn.Module):
    def __init__(self, size_img):
        super(Wrap, self).__init__()
        self.size_img = size_img

    def forward(self, imgs, boxes):
        outs = torchvision.ops.roi_align(imgs, boxes, 
                (self.size_img, self.size_img),
                spatial_scale=1.0, 
                sampling_ratio=-1)
        return outs


def main():
    model = Wrap(112)

    img_shape = (1, 3, 640, 640)
    boxes_shape = (1, 5)
    img_data = torch.randn(img_shape)
    boxes_data = torch.randn(boxes_shape)
    scripted_model = torch.jit.trace(model, (img_data, boxes_data)).eval()
    mod, params = relay.frontend.from_pytorch(scripted_model, {'input.1': img_shape, 'input.2':boxes_shape})

if __name__ == "__main__":
    main()

Does anyone have the same problem as me? Please write down your solution, thanks!

@alexwong @masahi

what is your torch and torchvision version?

Luckliy for you, if you are trying to convert roi align op, I have an example of doing just that in my PR https://github.com/apache/incubator-tvm/pull/4961.

3 Likes

Well, I took a few steps to solve my problem, thanks!

def convert_roi_align():
    def _impl(inputs, input_types):
        spatial_scale = inputs[2]
        pooled_size = (inputs[3], inputs[4])
        sampling_ratio = inputs[5]
        return tvm.relay.op.vision.roi_align(inputs[0], inputs[1],
            pooled_size, spatial_scale,
            sampling_ratio)
    return _impl
...
relay.frontend.from_pytorch(scripted_model, {'input.1': img_shape, 'boxes':boxes_shape}, {'torchvision::roi_align': convert_roi_align()})
...
1 Like