Added support for ensemble models

[stefdoerr]

Works fine but I don't like so much that the same function can return a model with 4 outputs (energy/force mean/std) or with 2 outputs (just energy, force) depending on if you pass a single file or multiple.
Maybe it could become an argument to load_model to disable the returning of the std?

[RaulPPelaez]

Depending on your final use of this you could return a Dict[str, Tensor] or a Tuple[Tuple[Tensor, Optional[Tensor]], Tuple[Tensor, Optional[Tensor]]] or a List[Dict[str, Optional[Tensor]]] storing mean and std for each output

[RaulPPelaez]

It would be interesting to make this class take *args,**kwargs for future proofing. I do not think it is compatible with TorchScript, tough.

[stefdoerr]

Yes I definitely considered that. But if you say it's not compatible we can leave it for now I guess

[RaulPPelaez]

Given that we are not aiming for TorchScript anymore, I would make this function variadic so that it is immune to changes in the TorchMD_Net interface.
This is particularly useful for when/if #306 is merged, which changes how we deal with charges, spins and such.

[RaulPPelaez]

I believe your implementation is clean enough for its purposes. Stacking the model outputs is wasteful, but the code is much cleaner and adaptable than a Welford-like an online computation.

Lets settle on the output and we can write a documentation entry for it and some tests. Maybe below this entry? https://torchmd-net.readthedocs.io/en/latest/models.html#loading-a-model-for-inference

[RaulPPelaez]

I added a sanity check:

def test_ensemble():
    ckpts = [join(dirname(dirname(__file__)), "tests", "example.ckpt")] * 3
    model = load_model(ckpts[0])
    ensemble_model = load_model(ckpts)
    z, pos, batch = create_example_batch(n_atoms=5)

    pred, deriv = model(z, pos, batch)
    pred_ensemble, deriv_ensemble, y_std, neg_dy_std = ensemble_model(z, pos, batch)

    torch.testing.assert_close(pred, pred_ensemble, atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(deriv, deriv_ensemble, atol=1e-5, rtol=1e-5)
    assert y_std.shape == pred.shape
    assert neg_dy_std.shape == deriv.shape
    assert (y_std == 0).all()
    assert (neg_dy_std == 0).all()

[stefdoerr]

Nice sanity test, thanks! Okay so we can merge it for now and we see in the future if it needs adaptation for different model types

[RaulPPelaez]

I added a section on this to the docs and made Ensemble variadic, just relaying the arguments to TorchMD_Net.forward.

This LGTM, feel free to merge if you agree.

[stefdoerr]

Looks lovely. Many thanks Raul!