I am trying to convert the FastPhotoStyle model (PyTorch) by NVIDIA. The stylization module takes four inputs i.e. the style image, the content image, and two segmentation masks.
I am running the conversion step like so -
torch.onnx.export(stylization_module, [cont_img, styl_img, cont_img, styl_img], 'stylization_module.onnx', input_names=['content', 'style', 'content_seg', 'style_seg'], output_names=['stylized_image'])`
It is producing:
RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type NoneType
Here’s the Colab Notebook for full reproducibility.
Any pointers to resolve the issue would be great.
It seems that the implementation of the
PhotoWCT model doesn’t use the standard
forward methods as seen in these lines of code, but instead apparently relies on the
The ONNX export however, uses the
forward method to call the model, which would return
None at the moment. I don’t know, why the model authors skipped these standard implementation and instead went with the custom
transform method (which would also break any data parallel approach etc.).
Anyway, you could monkey patch the model by assigning the
p_wct.forward and make sure the model executes properly in eager mode before trying to export it.
After I’ve patched the method, it seems that more errors are raised.
Thank you. So, if I understand correctly you did try the monkey patching way and got into errors, could you share them?
Sorry, don’t have the terminal open anymore, but just use
p_wct.transform(cont_img, styl_img, cont_img, styl_img) to raise this error.
Inside the export call, right? Sorry about the back and forth.
No, try to run the model without any exports first.
Once it’s working properly with the
transform method as well as with the direct model call (
p_wct(cont_img, ...)), which calls the
forward method, you could start trying to export it.
If your model isn’t working in eager mode first, the export would probably raise unhelpful errors.