.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated_examples/1-advanced/08-llpr-ensemble-training.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_1-advanced_08-llpr-ensemble-training.py: .. _llprensembleexample: Generating and training an LLPR-derived shallow ensemble model ============================================================== This tutorial demonstrates how to generate and, optionally, further train an LLPR-derived shallow ensemble model using metatrain. Building on the LLPR approach, this more advanced technique allows for 1) generation of a last-layer ensemble model from an LLPR model and 2) gradient-based tuning of ensemble weights using a negative log-likelihood (NLL) loss, often leading to improved uncertainty estimates at the cost of further training. We first train a baseline model without uncertainties: .. literalinclude:: options-model.yaml :language: yaml Then we create an LLPR ensemble model. This involves creating the LLPR model, which is very cheap (one pass through the training data without backpropagation), and then sampling last-layer ensemble weights using the LLPR covariance (extremely cheap), as explained in https://arxiv.org/html/2403.02251v1. Specifying `num_ensemble_members` enables the latter step in addition to the basic LLPR model. .. literalinclude:: options-llpr-ensemble.yaml :language: yaml In addition, you can decide to perform further backpropagation-based training on the resulting shallow ensemble, which is more expensive but can lead to better uncertainty estimates. This is done by setting `num_epochs` to the number of epochs that you want to train for. .. literalinclude:: options-llpr-ensemble-train.yaml :language: yaml You can train these models yourself with the following code: .. GENERATED FROM PYTHON SOURCE LINES 39-48 .. code-block:: Python import subprocess import ase.io import numpy as np from metatomic.torch import ModelOutput from metatomic.torch.ase_calculator import MetatomicCalculator .. GENERATED FROM PYTHON SOURCE LINES 49-51 We first train the baseline model without uncertainties, then train the LLPR ensemble models. .. GENERATED FROM PYTHON SOURCE LINES 52-70 .. code-block:: Python # Here, we run training as a subprocess. In practice, you would run this from # the command line, e.g., ``mtt train options-model.yaml -o model.pt``. print("Training baseline model...") subprocess.run(["mtt", "train", "options-model.yaml", "-o", "model.pt"], check=True) print("Training LLPR ensemble model...") subprocess.run( ["mtt", "train", "options-llpr-ensemble.yaml", "-o", "model-llpr-ens.pt"], check=True, ) print("Training LLPR ensemble model with further backpropagation...") subprocess.run( ["mtt", "train", "options-llpr-ensemble-train.yaml", "-o", "model-llpr-ens-tr.pt"], check=True, ) .. rst-class:: sphx-glr-script-out .. code-block:: none Training baseline model... Training LLPR ensemble model... Training LLPR ensemble model with further backpropagation... CompletedProcess(args=['mtt', 'train', 'options-llpr-ensemble-train.yaml', '-o', 'model-llpr-ens-tr.pt'], returncode=0) .. GENERATED FROM PYTHON SOURCE LINES 71-73 You can now use the uncertainties from the LLPR, as well as the ensemble model, as follows. .. GENERATED FROM PYTHON SOURCE LINES 74-108 .. code-block:: Python # Load some test structures structures = ase.io.read("ethanol_reduced_100.xyz", ":5") # Load the ensemble-trained model calc = MetatomicCalculator("model-llpr-ens.pt", extensions_directory="extensions/") # Get predictions with both ensemble and analytical uncertainties # (note that all these quantities are also available per-atom with ``per_atom=True``) predictions = calc.run_model( structures, { "energy": ModelOutput(per_atom=False), "energy_uncertainty": ModelOutput(per_atom=False), # LLPR analytical "energy_ensemble": ModelOutput(per_atom=False), # ensemble predictions }, ) energies = predictions["energy"].block().values.squeeze().cpu().numpy() llpr_uncertainties = ( predictions["energy_uncertainty"].block().values.squeeze().cpu().numpy() ) ensemble_predictions = ( predictions["energy_ensemble"].block().values.squeeze().cpu().numpy() ) # Calculate ensemble mean and standard deviation ensemble_mean = np.mean(ensemble_predictions, axis=1) ensemble_std = np.std(ensemble_predictions, axis=1) print(f"Energies: {energies}") print(f"LLPR analytical uncertainties: {llpr_uncertainties}") print(f"Ensemble mean: {ensemble_mean}") print(f"Ensemble std: {ensemble_std}") .. rst-class:: sphx-glr-script-out .. code-block:: none Energies: [-154.98422258 -154.98419238 -154.98523923 -154.98111695 -154.97973745] LLPR analytical uncertainties: [0.00816542 0.00742433 0.00558033 0.00600885 0.0057837 ] Ensemble mean: [-154.98422258 -154.98419238 -154.98523923 -154.98111695 -154.97973745] Ensemble std: [0.00774784 0.00755035 0.00529778 0.00613911 0.00533268] .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 35.874 seconds) .. _sphx_glr_download_generated_examples_1-advanced_08-llpr-ensemble-training.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 08-llpr-ensemble-training.ipynb <08-llpr-ensemble-training.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 08-llpr-ensemble-training.py <08-llpr-ensemble-training.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 08-llpr-ensemble-training.zip <08-llpr-ensemble-training.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_