How to write the SiC trained model to a KIM model in your example?

[ChesterCs-thu]

In your example of SiC model, how can I write the trained model_si and model_c to a KIM model for MD simulations ? Does the model.write_kim_model() function still work?
And when training multi-elemental material models, why should the models be separated? Can I train one model directly?

[mjwen]

Hi @ChesterCs-thu!

how can I write the trained model_si and model_c to a KIM model for MD simulations ? Does the model.write_kim_model() function still work?

Yes, you can still do model_si.write_kim_model() and model_c.write_kim_model(). This will generate two KIM models. But unfortunately, the current KIM NN model does not support multiple species, which means they cannot be readily used via KIM-API.

We are working on a more general way to interface Pytorch NN models with KIM. @ipcamit and @dskarls already have a working version, and we are testing it; should be available soon. Alternatively, we will have to modify the KIM NN model to support it.

why should the models be separated? Can I train one model directly?

The purpose for two models is to have two different parameter sets, one for Si and the other for C. Each NeuralNetwork only has one set of parameters. If we train one NeuralNetwork, there is currently no easy way to differentiate species.

[michaelmacisaac]

Good morning @mjwen ,
You mentioned you are working on a more general way to interface Pytorch NN models with KIM, does this include the ability to develop one model that includes two or more species?
Thanks

[ipcamit]

Hi @michaelmacisaac , yes. Currently we have really early stage version which can support almost any TorchScript model directly using a dedicated KIM model driver. Hoping to release it for alpha preview by the end of september, Idea is that any Pytorch model which accepts inputs in predefined fashion, and can be compiled to TorchScript, can be directly interfaced with KIM infrastructure. At present it lacks the support for Graph Neural Networks which we are working on.

For example, I have tried above SiC model in it by creating a torch.nn.Modulelist for Si and C then using that module list to evaluate C or Si environment.

[michaelmacisaac]

@ipcamit
Thanks! To clarify will this include functionality that will allow one model to be trained for multiple species?

[ipcamit]

Yes. This is my current implementation of multispecies SiC Model as of now.

N1 = 10
N2 = 10
model_si = Sequential(Linear(51, N1), Tanh(), Linear(N1, N2), Tanh(), Linear(N2, 1))
model_c = Sequential(Linear(51, N1), Tanh(), Linear(N1, N2), Tanh(), Linear(N2, 1))
class ModelSiC(torch.nn.Module):
    def __init__(self):
        super(ModelSiC,self).__init__()
        self.models = torch.nn.ModuleList([model_c,model_si])

    def forward(self,x:torch.Tensor, weight:torch.Tensor):
        energies = self.models[0](x) * torch.unsqueeze(weight[:,0],1) + self.models[1](x) * torch.unsqueeze(weight[:,1],1)
        return energies

model = ModelSiC()
model_jit = torch.jit.script(model)

Here weight is an 1-hot element vector, for C in above case it is [1,0], for Si it is [0,1]. It was needed in above format to ensure that module can compile to TorchScript. As of now it is bit wasteful as you have to compute the NN twice in each run. But once we move from development to model demonstrations, we would try to come up with a better solution.
For now this "works" and I can see it training as well.

[michaelmacisaac]

@ipcamit
Thanks for the help and the inclusion of your code.

[michaelmacisaac]

Yes, you can still do model_si.write_kim_model() and model_c.write_kim_model(). This will generate two KIM models. But unfortunately, the current KIM NN model does not support multiple species, which means they cannot be readily used via KIM-API.

@mjwen
To clarify, are you saying that if we were to write two models, one for Si and one for C, we could not perform a single MD study using the two models?
I'm unsure what the current capabilities of the KIM-API are and whether I can train two models for a two-element system and implement said models in a single MD study. Please clarify

[mjwen]

KIM-API is a general package for interfacing potential models with simulation code such as LAMMPS. Within KIM-API, one has to implement specific potential models. The current KIM model for NN potentials only support single species. It needs to be expanded to support multiple species.

So, yes, if you save two models via kliff, the current KIM model does not know how to handle it. We are working on this and will have a general KIM model supporting multiple species.