Use of djl with AIMNet2

[mjw99]

I am currently trying to interface with a PyTorch model, AIMNet2 ( https://github.com/isayevlab/AIMNet2/ ). I am using djl version 0.24.0 with the following code:

package name.mjw.djl;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;

import ai.djl.MalformedModelException;
import ai.djl.Model;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;

import ai.djl.training.util.DownloadUtils;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.Batchifier;
import ai.djl.translate.TranslateException;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;

/**
 * Example based upon:
 * 
 * python ./calculators/aimnet2_ase_opt.py ./models/aimnet2_wb97m-d3_ens.jpt ./water.xyz water_min.xyz
 */
public class AIMNet2 {

	public static void main(String[] args)
			throws MalformedModelException, IOException, NumberFormatException, TranslateException {

		DownloadUtils.download(
				"https://raw.githubusercontent.com/isayevlab/AIMNet2/main/models/aimnet2_wb97m-d3_ens.jpt",
				"aimnet2_wb97m-d3_ens.jpt", new ProgressBar());

		Path modelDir = Paths.get("./");
		Model model = Model.newInstance("aimnet2_wb97m-d3_ens.jpt");
		model.load(modelDir);

		System.out.println(model);

		Translator<NDList, NDList> translator = new Translator<NDList, NDList>() {

			@Override
			public NDList processInput(TranslatorContext ctx, NDList input) {
				return input;
			}

			@Override
			public NDList processOutput(TranslatorContext ctx, NDList output) {
				return output;
			}

			@Override
			public Batchifier getBatchifier() {
				return Batchifier.STACK;
			}
		};

		Predictor<NDList, NDList> predictor = model.newPredictor(translator);

		NDManager manager = NDManager.newBaseManager();

		// See
		// http://djl.ai/docs/faq.html#how-can-i-pass-arbitrary-input-data-type-to-a-pytorch-model
		NDArray coord = manager.create(
				new float[] {
						 0.7494f, 0.0000f, 0.4424f,
						 0.0000f,  0.0000f, -0.1654f,
						-0.7494f,  0.0000f,  0.4424f },
				new Shape(1, 3, 3));
		coord.setName("input1.coord");

		System.out.println(coord);

		NDArray numbers = manager.create(new long[] { 1, 8, 1 }, new Shape(1, 3));
		numbers.setName("input1.numbers");
		System.out.println(numbers);

		NDArray charge = manager.create(new float[] { 0.f }, new Shape(1));
		charge.setName("input1.charge");
		System.out.println(charge);

		NDList input = new NDList(coord, numbers, charge);

		System.out.println(input);

		System.out.println(predictor.predict(input));

	}
}

When running this, I see:

<SNIP>
  File "/data/roman/AIMNet2Paper/models/d3bj.py", line 81, in forward
    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        c6ij, d_ij = self._calc_c6ij(data)
                     ~~~~~~~~~~~~~~~ <--- HERE
        r4r2 = self.r4r2[data['numbers']]
        rrij = 3 * r4r2.unsqueeze(-1) * r4r2.unsqueeze(-2)
  File "/data/roman/AIMNet2Paper/models/d3bj.py", line 72, in _calc_c6ij
        w = l_ij.flatten(-2, -1).sum(-1)
        # multiply-sum is better written as matrix multiplication
        z = torch.einsum('bnmij,bnmij->bnm', c6ref, l_ij)
            ~~~~~~~~~~~~ <--- HERE
        _w = w < 1e-5
        z[_w] = 0.0
RuntimeError: einsum(): the number of subscripts in the equation (5) does not match the number of dimensions (6) for operand 0 and no ellipsis was given

	at ai.djl.pytorch.jni.PyTorchLibrary.moduleRunMethod(Native Method)
	at ai.djl.pytorch.jni.IValueUtils.forward(IValueUtils.java:57)
	at ai.djl.pytorch.engine.PtSymbolBlock.forwardInternal(PtSymbolBlock.java:145)
	at ai.djl.nn.AbstractBaseBlock.forward(AbstractBaseBlock.java:79)
	at ai.djl.nn.Block.forward(Block.java:127)
	at ai.djl.inference.Predictor.predictInternal(Predictor.java:144)
	at ai.djl.inference.Predictor.batchPredict(Predictor.java:184)
	... 2 more

Are there any clues as to what I am doing wrong?

Many thanks,

Mark

[frankfliu]

@mjw99
The code looks fine. Can you make sure the input shape are expected? It seems you are missing one dimension. Do you have a simple python code that creates the same input and just run forward for your model?

[mjw99]

Dear Frank,

Thank you for looking at this. The following protocol will install AIMNet2 using Mamba and carry out an optimisation on a molecule of water:

mamba create -n aimnet2 python==3.11 -y

mamba activate aimnet2

mamba install ase openbabel -y
pip3 install torch==2.0.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

git clone https://github.com/isayevlab/AIMNet2.git

cd AIMNet2

cat <<EOF > water.xyz 
3
Water
H    0.7493682    0.0000000    0.4424329
O    0.0000000    0.0000000   -0.1653507
H   -0.7493682    0.0000000    0.4424329
EOF

python ./calculators/aimnet2_ase_opt.py  ./models/aimnet2_wb97m-d3_ens.jpt  water.xyz water_min.xyz

This should yield:

Loading AIMNet2 model from file ./models/aimnet2_wb97m-d3_ens.jpt
       Step     Time          Energy         fmax
LBFGS:    0 09:14:42    -2081.040424        0.6370
LBFGS:    1 09:14:42    -2081.044630        0.2391
LBFGS:    2 09:14:43    -2081.046007        0.1633
LBFGS:    3 09:14:43    -2081.047413        0.0033

As an aside, the code responsible for packaging up the data to the model can be found here: https://github.com/isayevlab/AIMNet2/blob/3ce06f4002eb35a5bf4a2b9daf8b4f09be9aabb8/calculators/aimnet2ase.py#L34

[mjw99]

Also seeing this with djl 0.25.0

[mjw99]

Also, tried to force the use of pytorch 2.0 via:

                <dependency>
                    <groupId>ai.djl.pytorch</groupId>
                    <artifactId>pytorch-native-cpu</artifactId>
                    <classifier>linux-x86_64</classifier>
                    <version>2.0.1</version>
                    <scope>runtime</scope>
                </dependency>

                <dependency>
                    <groupId>ai.djl.pytorch</groupId>
                    <artifactId>pytorch-jni</artifactId>
                    <version>2.0.1-0.25.0</version>
                    <scope>runtime</scope>
                </dependency>

.....
[name.mjw.djl.AIMNet2.main()] INFO ai.djl.util.Platform - Found matching platform from: jar:file:/home/mjw/.m2/repository/ai/djl/pytorch/pytorch-native-cpu/2.0.1/pytorch-native-cpu-2.0.1-linux-x86_64.jar!/native/lib/pytorch.properties

But, again no joy.

[mjw99]

Also, running the reference python implementation on the CPU

pip3 install torch==2.0.1+cpu -f https://download.pytorch.org/whl/torch_stable.html

works fine.

[mjw99]

An updated minimal standalone Python example that can be used with just PyTorch and the model's jpt file:

import torch

# https://raw.githubusercontent.com/isayevlab/AIMNet2/main/models/aimnet2_wb97m-d3_ens.jpt
model = torch.jit.load('aimnet2_wb97m-d3_ens.jpt')

# Water molecule: H-O-H
coord = torch.as_tensor([   (0.7494, 0.0000,  0.4424),
                            (0.0000, 0.0000, -0.1654), 
                            (-0.7494, 0.0000,  0.4424)]).to(torch.float).unsqueeze(0)
numbers = torch.as_tensor([1, 8, 1]).to(torch.long).unsqueeze(0)
charge = torch.as_tensor([0.0]).to(torch.float)

d = dict(coord=coord, numbers=numbers, charge=charge)

out = model(d)

ret = dict(energy=out['energy'].item())
print(ret)

[frankfliu]

@mjw99

I tried your model (https://raw.githubusercontent.com/isayevlab/AIMNet2/main/models/aimnet2_wb97m-d3_ens.jpt), and I got different error:

  File "code/__torch__/ensemble/___torch_mangle_4.py", line 27, in forward
      else:
        pass
    _out = (_0).forward(_in, )
            ~~~~~~~~~~~ <--- HERE
    _r = annotate(Dict[str, Tensor], {})
    _6 = torch.keys(_out)
  File "code/__torch__/aimnet/modules.py", line 34, in forward
    _4 = [data0[x0]]
    training = self.training
    _5 = torch.autograd.grad(_3, _4, None, None, training)
         ~~~~~~~~~~~~~~~~~~~ <--- HERE
    g = _5[0]
    if torch.__isnot__(g, None):

If I turn on inference model with python code I got the same error:

    with torch.inference_mode():
        out = model(d)

Looks like when you trace your model, it's run in training mode. In DJL, we assume, you run inference model, and autograd is turn off in DJL. Can you turn off training model when you trace the model?

[mjw99]

Dear Frankfliu,
Apologies for the delay in replying to this.

Can you turn off training model when you trace the model?

I am not sure all of the model is available on the AIMNet2 repo, hence I am not sure how this can be done. But, thank you for looking at this; it is appreciated.