Jittable trained model

[mprithiv]

Based on this link and this one, I am trying to convert my GraphSAGE model to a JIT traced version. I see the model is converted to JIT at line num 35 in the gat.py

model = torch.jit.script(model)

And then, the train function is defined and the model is trained. Does the model forward pass needs dummy inputs to be specified?

class PyTorch_to_TorchScript(torch.nn.Module): def __init__(self): super(PyTorch_to_TorchScript, self).__init__() self.model = model def forward(self, data, edge_index_0, edge_size_0, edge_index_1, edge_size_1, edge_index_2, edge_size_2): return self.model(data, edge_index_0, edge_size_0, edge_index_1, edge_size_1, edge_index_2, edge_size_2)

[rusty1s]

You will only need dummy inputs for conversion via torch.jit.trace. torch.jit.script relies on typed arguments and TorchScript and does not need any dummy inputs.

[mprithiv]

Is there any example of how to create a client side (triton) interaction with JIT script GNN models.

[rusty1s]

How about:

• https://medium.com/nvidia-ai/how-to-deploy-almost-any-pytorch-geometric-model-on-nvidias-triton-inference-server-with-an-218d0c0c679c
• https://github.com/pyg-team/pytorch_geometric/tree/master/examples/cpp

[mprithiv]

is it mandatory to create dummy inputs and outputs for every layer as mentioned in the first link? What would be the case for a heterograph with HGTloader

[rusty1s]

I actually don't know since I have no experience with Triton. As far as I am aware, it requires static sizes across different batches. With HGTLoader, you can achieve a fixed size of nodes, but the number of edges may still vary.