Questions for porting GNN models

[mjwen]

Hi @ipcamit Thanks for this great work!

The major comment I have is it would be great if you could create a step-by-step guide to convert some simple models to be compatible with the driver. I think such a tutorial would make it much easier for users.

The model can be a 2-layer gnn (does not need to worry about smoothness and such, just as simple as possible) that is fit to two species. Also, it would be great if the model uses supercell vectors and minimal image convention to deal with PBCs, because most of the GNN models do that.

I was imagining something like an example directory, labeled by 1-xxx, 2-xxx, 3-xxx..., where in each directory the input and output are provided, as well as the commands/steps to convert the input to the output. Then a user like me can work through it to get more familiar with it before we start to port our models.

I have a couple of questions:

1. Will a function/module decorated by torch.jit.trace work?

2. Why the CPP version of PyG is needed in the driver; can we remove this dependence? Otherwise, DGL models cannot be updated?

3. In torchmlAPI, does MY_TORCH_MODEL.pt and MY_TORCH_MODEL.pt mean the same file?

4. In gnn.md, what do you mean by box_dims. I cannot think of existing gnn models using that.

5. should the cutoff be influence distance below?

   

6. It is difficult for me to understand this snippet below and the docs around that. Can you expand on it? Also, should the h_i = phi_h... line be moved out of the else statement?

   

7. Appendix A1 might need to be reworked? I understand the purpose is to demonstrate how to use a wrapper model, but a more fleshed example would be more helpful. Two quick questions: first, the two models have the same signature; why bother using the WrapperModel? second, why not directly call the forward method of MyModel in WrapperModel, but call the children modules in the model? I think a different example might be needed.

[ipcamit]

Thanks for the comments. I think the example idea is great, and I will add then ASAP. In the mean let try and answer the remaining questions:

1. Will a function/module decorated by torch.jit.trace work?

Yes. But you would still need to save the module/function as func.save("func.pt") etc. Actually, this is a great idea to give example of a simple LJ potential function as a first step.

2. Why the CPP version of PyG is needed in the driver; can we remove this dependence? Otherwise, DGL models cannot be updated?

No we cannot. Sadly DGL is not compatible with the TorchML model driver. Major reason being, last I checked, DGL does not provide jittable layers like PyG does. Hence DGL cannot be used with C++ interface of PyTorch.

3. In torchmlAPI, does MY_TORCH_MODEL.pt and MY_TORCH_MODEL.pt mean the same file?

I think you meant MY_TORCH_MODEL.pt and MyTorchScriptFile.pt? This is a typo, will rename MY_TORCH_MODEL.pt to MyTorchScriptFile.pt.

4. In gnn.md, what do you mean by box_dims. I cannot think of existing gnn models using that.

I think I would need to explain it better, what I meant is this: Line 71 onwards in nequip. It used cell vectors in the computation. In LAMMPS etc equivalent would be box bims (xlo, xhigh etc, see Line 325 )

should the cutoff be influence distance below?

No this is cutoff only, for each atom you now only need the cutoff, as all the atoms have neighbors now (upto infl distance of center atom).

6. It is difficult for me to understand this snippet below and the docs around that. Can you expand on it? Also, should the h_i = phi_h... line be moved out of the else statement?

Yes, it should be outside. will edit. Will add more explanation or better code snippet. I just copy pasted from an older GNN I had.

why bother using the WrapperModel? second, why not directly call the forward method of MyModel in WrapperModel, but call the children modules in the model? I think a different example might be needed.

Good idea, ideally you do not need the wrapper here. I guess it is not a good example that way. Let me add a better one where you need the wrapper.