Using jit to convert VQA model to iPhone (MMF)

I’m trying to convert a state_dict Visual Question Answering “.pth”-model using jit.script.

The model has been made through facebooks MMF-framework, and takes as input a picture, a question as well as utilises a Knowledge Graph.

While scripting the model, I’ve encountered various problems, which I’ve continuously tried to fix. However, as this is a massive framework, utilising various scripts as well as native python scripts, the quest to script the model seems endless…

Now, I’ve arrived at a problem that I’m not sure is fixable through Jit, and I’m asking for help.

Jit.script has encountered a function, which takes another function as input. This function can be various ‘Forward’-functions (if anyone is familiar with MMF, if not, you don’t have to understand this part), which means I can’t just manually implement the input-function as part of the scripted function as I’ve had to do previously.

As Jit needs type-hinting, I’ve had to annotate the input as a “Callable[…, Tensor]” type. However, Jit doesn’t support this annotation-type. Is there any way I can work my way around this?

I’ve also had to resort to deleting various Args and Kwargs in native python files such as, ruining my Python in the process, in the hopes of scripting this model. However, that seems like the wrong course of action.

I feel lost, is scripting this model a lost cause or do you know of any way I can bypass these problems?

One possibility is to use torch.jit.interface pytorch/ at master · pytorch/pytorch · GitHub