Inference with dictionary of pytorch tensors

[jonathan-booth]

I'm trying to use FTorch to call a pytorch model from fortran 90. The pytorch model wants a dictionary of torch tensors, for example, inference in python looks like this:

`pos=torch.tensor([[1.0,1.1,1.2],
[0.0,0.1,0.2],
[2.1,2.0,2.2],
])

types=torch.tensor([0,1,0])

edges=torch.tensor([[0, 0, 1, 1, 2, 2],
[1, 2, 0, 2, 0, 1]])

input={
'pos':pos,
'edge_index':edges,
'atom_types': types
}

output=model(input)`

Can I do this with FTorch? There seem to be various workarounds for building dictionaries or dictionary-like objects in fortran 90.

Thanks

[jatkinson1000]

Hi @jonathan-booth At the moment if you have a model that takes multiple input tensors we facilitate this by providing them as individual items rather than as a single dictionary of tensors.

You can see an example of this in our Multiple inputs and outputs worked example which takes two input tensors (they don't need to be the same shape or size). In python this is declared as two separate inputs to the forward method, and in Fortran this is done using an array of two torch_tensors.

So my advice, unless you are for some reason restricted to using a dict, would be to write your net architecture as taking each input tensor as a distinct input. This would then be possible to can using FTorch as in the above example.

If this helps please let me know, and if not we would be happy to take a closer look at some code.

[jonathan-booth]

Thanks. I'm calling a model made by the allegro package:

https://github.com/mir-group/allegro

so I don't really have control over how it takes its input, although maybe I could modify it so it takes its input as an array of tensors like in the example you gave. If this doesn't prove to be possible is there a way to get FTorch to pass a dict of torch tensors to the model?

Thanks

[jatkinson1000]

@jonathan-booth Thanks, that helps.

Looking at the allegro code they have a lot of infrastructure build around a Torch neural net, so you would probably have to do some surgery to extract the net so that you could save it as TorchScript. (I'm looking at https://github.com/mir-group/allegro/blob/main/allegro/nn/_fc.py) If you do this then I would also take the opportunity to modify so that you pass inputs in as torch tensors.

The beauty of FTorch is that it bypasses the python runtime to use the more efficient C++ implementation of Torch. However, it is not clear to me whether you can save the allegro model to TorchScript or not. If this is not possible you might need to look at the more 'traditional' methods for coupling python to Fortran.

[TomMelt]

An idea (disclaimer not sure if it's possible) but would it be easier to:

• save the model in another more python friendly format i.e., save the architecture and weights of the model.
• load said model and tweak in python
• save the tweaked model to torchscript
• call from Fortran using Ftorch

🤷

[jatkinson1000]

@jonathan-booth Oh, that's great! In that case if you can change the input it takes we should be away.

Looking at the C++ API suggests it may be possible to support passing input tensors as a dictionary though s subtype of IValue: https://pytorch.org/cppdocs/search.html?q=dict&check_keywords=yes&area=default

However, we don't support this at present so you would need to modify the c++ part of the code as well as the Fortran in src/.
We'd be happy for you to experiment and add this functionality though!

[ElliottKasoar]

It sounds like saving to TorchScript is already solved, but just to clarify in case this isn't what you already figured out, @jonathan-booth: my understanding is that Allegro workflows (e.g. tutorial) tend to make use of NequIP commands, including nequip-deploy, so I think most tooling is built around saving trained models to TorchScript already.

I agree with @TomMelt that loading and modifying the model is probably the best immediate solution, if it's possible. It's not something I've ever done, but I would hope that even if you can't modify the existing layers, it might be possible to prepend an additional layer that could convert from a list to the required dictionary.

This avoids changing the Allegro source code, which risks conflicts if further changes are made. NequIP certainly warns that the code is under active development and formats/APIs may change, although Allegro appears relatively inactive, and given that a few other models for interatomic potentials seem to outperform it, perhaps this will remain the case.

[jonathan-booth]

The method of loading the model and modifying the input has worked. I'm now trying to modify the output as well because the output is a Dict[str: tensor]. This is proving difficult. Does FTorch work with models that have been scripted with torch.jit.script? I think that's the only way I can save my modified model. I do the following:

class ModifiedModel(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.original_model = original_model
    
    def forward(self, new_input):
        processed_input = self.change_input(new_input)
        old_output = self.original_model(processed_input)
        new_output=self.change_output(old_output)

        return new_output
    
    def change_input(self, old_input):
        # processed_input = some modification of old_input
        return processed_input
 
    
    def change_output(self, old_output: Dict[str, torch.Tensor]):
        # processed_output = some modification of old_output
        # processed_output is an array of floats.
        return processed_output

original_model= torch.jit.load('old_model.pt')
modified_model = ModifiedModel(original_model)
scripted_model = torch.jit.script(modified_model)
torch.jit.save(scripted_model,'new_model.pt')

If FTorch is fine with this then the problem must be with my understanding of torchscript.

[jatkinson1000]

Yes, scripting or tracing should be fine.

[jonathan-booth]

OK it all works now, the issue with the output was my error. Thanks a lot for your help!

[jatkinson1000]

@jonathan-booth That's great to hear, thank you.

If you could provide a short summary of how you resolved the issue for future readers that would be really helpful.
If the answer is in one of the comments above please could you "mark it as answer" to close the discussion?
Thanks.

[jonathan-booth]

OK will do. The error in changing the output was just a trivial typo in my change_output function. Therefore, the comment I made above with my ModifiedModel class is enough to do what I ask in this discussion. I'll mark that as the answer. Thanks for your help!