JIT-ing Question

[CrossR]

Hey!

Thanks for the library, its been a big help so far.

I'm now looking into how easy it is to move a model from Python to C++ via Torchscript.
I've looked through it for basic pytorch, then the bits in the docs for pytorch_geometric.

I was testing it on the ogbn_proteins_deepgcn.py example, as it uses some of the same modules I use.
I stuck .jittable() on the two geo bits, but receieved an error when trying to apply that to DeepGCNLayer (saying it doesn't have jittable), so I dropped it there thinking maybe since its a wrapper bit it doesn't need that.

(Code in a Gist for completeness. I dropped the dataset loading since I figured it wasn't going to change all that much, and I'd need to install another package into my venv.)

However I instead get this issue:

raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/Users/rcross/.virtualenvs/uni-work-x0bl1lnP-py3.8/lib/python3.8/site-packages/torch_geometric/nn/models/deepgcn.py", line 67
    def forward(self, *args, **kwargs):
                              ~~~~~~~ <--- HERE
        """"""
        args = list(args)

My (naive) assumption here would be I should just make my own "copy" of deepgcn and swap to using non-keyword arguments. Is that the best process here, or am I doing something wrong? Its entirely possible I've misunderstood the documentation and I'm missing a step/part of the process is missing at the moment.

Also, if you have any examples/documentation of building a graph or similar in C++, that would be appreciated!
I understand the process of Python Model -> Torchscript Model -> Load Model, but I've struggled to find how that model will then expect my data to look in C++.

I'd be happy to make a PR later on with an extended torchscript example, assuming I get this working.

[wsad1]

I think you are right, implementing a version of DeepGCNLayer where the forward has arguments with type hints should work, as described here. Also you would need to call jittable on the convolution layers inside the DeepGCNLayer as described here.

[CrossR]

I gave it a try, and edited DeepGCNLayer to drop the keyword args (i.e. swap it directly to def forward(self, x, edge_index, edge_attr):), and propagated that change around to the few other places using args and kwargs and that seemed to solve the issue.

I hit a few other issues that were fine to get past (Can't use a layer directly, need to swap to using layer.forward, can't pass data and decompose it, should just pass the elements around directly), but I couldn't work around an issue with checkpoint.

I can of course, just remove it, but its potentially going to be an issue for someone else in the future. (If anyone is wondering, you get a keyword arg error again, and if you try and avoid that, it complains about <built-in method apply of FunctionMeta object at 0x7f8a635c3920> is currently not supported in Torchscript).

My model at least loads into my example app now, so just need to work out what tensor shape it is expecting.

[chrisba73]

My model at least loads into my example app now

@CrossR What did you have to do to load your Python JIT model into C++? I'm not having any luck, and I'm now wondering if/how this is supported? For now, I'm trying to just do the simple model here, converting with:

model = torch.jit.script(model)
model.save('gcn_jit.pt')

and I'm getting this error in C++ when trying to read in the JIT model [first few lines shown]:

error loading the model: schemas.size() > 0INTERNAL ASSERT FAILED at "..\\..\\torch\\csrc\\jit\\frontend\\schema_matching.cpp":525, please report a bug to PyTorch.
Exception raised from matchSchemas at ..\..\torch\csrc\jit\frontend\schema_matching.cpp:525 (most recent call first):
00007FF9766A051A00007FF97669F5E0 c10.dll!c10::detail::LogAPIUsageFakeReturn [<unknown file> @ <unknown line number>]
00007FF9766A007A00007FF97669F5E0 c10.dll!c10::detail::LogAPIUsageFakeReturn [<unknown file> @ <unknown line number>]
00007FF9766A127100007FF97669F5E0 c10.dll!c10::detail::LogAPIUsageFakeReturn [<unknown file> @ <unknown line number>]
00007FF9766A0EE500007FF97669F5E0 c10.dll!c10::detail::LogAPIUsageFakeReturn [<unknown file> @ <unknown line number>]
00007FF97669EAEF00007FF97669EA80 c10.dll!c10::Error::Error [<unknown file> @ <unknown line number>]
00007FF97669D40600007FF97669D340 c10.dll!c10::detail::torchCheckFail [<unknown file> @ <unknown line number>]
00007FF97661D4E600007FF97661D4A0 c10.dll!c10::detail::torchInternalAssertFail [<unknown file> @ <unknown line number>]
00007FF95BCB22DC00007FF95BCB2240 torch_cpu.dll!torch::jit::matchSchemas [<unknown file> @ <unknown line number>]
00007FF95BCB2F1100007FF95BCB2A50 torch_cpu.dll!torch::jit::emitBuiltinCall [<unknown file> @ <unknown line number>]
00007FF95BCCD0C900007FF95BCCCFC0 torch_cpu.dll!torch::jit::BuiltinFunction::call [<unknown file> @ <unknown line number>]
00007FF95BC6BFAA00007FF95BC49130 torch_cpu.dll!torch::jit::ScriptTypeParser::operator= [<unknown file> @ <unknown line number>]
00007FF95BC89F7500007FF95BC49130 torch_cpu.dll!torch::jit::ScriptTypeParser::operator= [<unknown file> @ <unknown line number>]
00007FF95BC85EBA00007FF95BC49130 torch_cpu.dll!torch::jit::ScriptTypeParser::operator= [<unknown file> @ <unknown line number>]
00007FF95BC7298C00007FF95BC49130 torch_cpu.dll!torch::jit::ScriptTypeParser::operator= [<unknown file> @ <unknown line number>]
...

with code that is simply doing this (which I've used fine on other non-Pytorch Geometric JIT models, it works):

    torch::jit::script::Module module;
    try {
        module = torch::jit::load(path_to_model);
    }
    catch (const c10::Error& e) {
       std::cerr << "error loading the model: " << e.what() << std::endl;
    }

[CrossR]

I didn't hit any issues on the C++ side, outside of needing to compile and build against a few of the related libraries, such as pytorch_sparse etc, and include their headers in my C++ code. The actual code I was using in C++ looks identical to yours (taken from one of the tutorials on the pytorch site).

Missing libraries gave very clear errors though, about missing operations so it doesn't look like that is the case here.

From a quick look around online, the first recommendation seems to be ensuring the versions of torch are consistent (for both Python and C++).

There is also this issue: pytorch/pytorch#35881 which implies it may be a Windows only issue (I'm on Mac/Linux).

[chrisba73]

Thanks for the quick reply @CrossR, I appreciate your feedback! (Yes, I'm on Windows; I don't think I have any DLL issues, nvidia-smi, conda, my .props files in MSVC, and depends.exe(my EXEs and DLLs) are all pointing to things that are CUDA 11.1.).

In Visual Studio, hitting F11 repeatedly on the torch::jit::load call, and giving it current libtorch and windows files for source when asked, I was able to get to a point in schema_matching.cpp, emitBuiltinCall(), where an error is built and thrown with the message: "Could not find any similar ops to ptr2ind. This op may not exist or may not be currently supported in Torchscript". But this error is different from the one I get from stepping over with "F10", which I posted above.

AHA: just now figured out why. emitBuiltinCall() throws torch::jit::ErrorReport [uncaught], then calls matchSchemas(), which throws TORCH_INTERNAL_ASSERT. So, this code

    torch::jit::script::Module module;
    try {
        module = torch::jit::load(path_ts_td_model);
    }
    catch (const torch::jit::ErrorReport& e) {
        std::cerr << "\nErrorReport: " << e.what() << std::endl;
    }
    catch (const c10::Error& e) {
        std::cerr << "\nGeneral error loading the model: " << e.what() << std::endl;
    }

now gives a much more specific (and beautifully drawn) error report which points to torch.ops.torch_sparse.ptr2ind. OK, great.

So, now trying to build pytorch_sparse via CMake. For FindTorch.cmake I gave it the version in the debug libtorch. CMake gave several warnings related to linking to Python3::Python and empty CUDA_ARCHITECTURES, but it said it was for project developers so...since I saw a torchsparse.sln file, I pressed on. But, when trying to compile, I now get tons of CUDA compile errors with just "exited with code 1". The error messages give the offending nvcc command; if I run the same command in the x64 Native tools window, I get "A single input file is required for a non-link phase when an outputfile is specified". I can get past this error message by deleting the command line switches like /Z7 /EHsc /wd4267 ... but then I get an error that it can't find Python.h. I can add an include to my Anaconda 3.9.5 directory, and it compiles, but gives an error at the end:

.../AppData/Local/Temp/tmpxft_00000b14_00000000-7_convert_cuda.cudafe1.cpp: fatal error C1041: cannot open program database 'C:\...\pytorch_sparse\csrc\cuda\torchsparse.dir\Debug\vc142.pdb'; if multiple CL.EXE write to the same .PDB file, please use /FS

But the command line does have the /FS switch. I'm thinking that maybe the command line switch format of the nvcc command that's being formed is a little confusing, maybe because of some CMake/Windows combination (supported a little by this).

Hopefully I can get a pytorch_sparse built on Windows!