prim::PythonOp is not supported when convert GPT2 jit trace to relay IR

Hi I encounter following error when convert JIT trace to tvm.relayIR:
NotImplementedError: The following operators are not implemented: [‘prim::PythonOp’, ‘aten::ger’, ‘aten::or’, ‘aten::view_as’, ‘aten::index_put_’]

I know how to handle other four operators except prim::PythonOp. Does anyone know how to solve this unimplemented Operator?
If I print out jit graph, I can’t saw prim::PythonOp in the graph. Thus I don’t know where it comes from. I definitely use model.eval() and I know there’s post about unimplemented PythonOp which can be solved with tvm.relay.op.vision.roi_align. But my issue here is the model is very big and I don’t know where does this PythonOp come from.

Any suggestions?

1 Like

PythonOp is an arbitrary op written in python, there is no way we can convert from it.

@masahi Thanks for your quick response. Thus the possible solution should be: First, find out which code logic use PythonOp like roi_align case. Then, we check whether we can use tvm.relay.op.xxx to replace that logic. If can, we replace tvm op with that logic in model and re-gen jit trace and convert to relay IR. If can’t, then we can’t convert this jit graph by tvm.from_pytorch directly. Am I right?

What GPT2 you are using? I was using the one from transformers 3.5 without issues.

2 Likes

Yes, if you can map your python code to corresponding combination of relay ops, it could work. But this is always application specific so there is no generic way to support PythonOp in our “official” frontend. You can use the custom convert map to register your conversion.

1 Like

I met similar error when convert GPT jit trace to relay IR. Could you please have a look at it or share your demo code? Thanks in advance~

The error is:

ValueError: value has to be scalar or NDArray`

And the script used is as following:

import torch
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer


TRACED_PT_FILE_NAME = "traced_gpt2.pt"


# 1. Convert to TorchScripted model
token_predictor = GPT2LMHeadModel.from_pretrained("gpt2", torchscript=True).eval()

# trace
random_tokens = torch.randint(10000, (5,))
traced_token_predictor = torch.jit.trace(token_predictor, random_tokens)
torch.jit.save(traced_token_predictor, TRACED_PT_FILE_NAME)


# 2. Use PyTorch inference
sentence_fragment = "The Manhattan bridge is a major"

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
context = torch.tensor(tokenizer.encode(sentence_fragment))

loaded_model = torch.jit.load(TRACED_PT_FILE_NAME).eval()
# print(loaded_model.graph)
torch_out = loaded_model(context)


# 3. Use TVM
import tvm
from tvm import relay

inputs = [("dummy_input_name", (5,))]
mod, params = relay.frontend.from_pytorch(loaded_model, inputs)
print(mod)