deploy: 80cf309

_downloads/f926bfac3d372dde77d6e6dc68532a9c/embedding-monkeypatch.py

@@ -2,6 +2,7 @@
 
 import torch
 from ocpmodels.common.utils import conditional_grad, scatter_det
+from ocpmodels.models.gemnet_oc.utils import repeat_blocks
 
 
 @conditional_grad(torch.enable_grad())

_sources/autoapi/ocpmodels/common/registry/index.rst

@@ -123,8 +123,8 @@ Attributes
 
           from ocpmodels.common.registry import registry
 
-          @registry.register_logger("tensorboard")
-          class WandB():
+          @registry.register_logger("wandb")
+          class WandBLogger():
               ...
 
 

_sources/autoapi/ocpmodels/models/equiformer_v2/trainers/energy_trainer/index.rst

@@ -25,7 +25,7 @@ Classes
 
 
 
-.. py:class:: EquiformerV2EnergyTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='tensorboard', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
+.. py:class:: EquiformerV2EnergyTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='wandb', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
 
 
    Bases: :py:obj:`ocpmodels.trainers.OCPTrainer`
@@ -68,7 +68,7 @@ Classes
                 (default: :obj:`None`)
    :type seed: int, optional
    :param logger: Type of logger to be used.
-                  (default: :obj:`tensorboard`)
+                  (default: :obj:`wandb`)
    :type logger: str, optional
    :param local_rank: Local rank of the process, only applicable for distributed training.
                       (default: :obj:`0`)

_sources/autoapi/ocpmodels/models/equiformer_v2/trainers/forces_trainer/index.rst

@@ -25,7 +25,7 @@ Classes
 
 
 
-.. py:class:: EquiformerV2ForcesTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='tensorboard', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
+.. py:class:: EquiformerV2ForcesTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='wandb', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
 
 
    Bases: :py:obj:`ocpmodels.trainers.OCPTrainer`
@@ -68,7 +68,7 @@ Classes
                 (default: :obj:`None`)
    :type seed: int, optional
    :param logger: Type of logger to be used.
-                  (default: :obj:`tensorboard`)
+                  (default: :obj:`wandb`)
    :type logger: str, optional
    :param local_rank: Local rank of the process, only applicable for distributed training.
                       (default: :obj:`0`)

_sources/autoapi/ocpmodels/trainers/base_trainer/index.rst

@@ -25,7 +25,7 @@ Classes
 
 
 
-.. py:class:: BaseTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier: str, timestamp_id: Optional[str] = None, run_dir: Optional[str] = None, is_debug: bool = False, print_every: int = 100, seed: Optional[int] = None, logger: str = 'tensorboard', local_rank: int = 0, amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm={}, noddp: bool = False)
+.. py:class:: BaseTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier: str, timestamp_id: Optional[str] = None, run_dir: Optional[str] = None, is_debug: bool = False, print_every: int = 100, seed: Optional[int] = None, logger: str = 'wandb', local_rank: int = 0, amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm={}, noddp: bool = False)
 
 
    Bases: :py:obj:`abc.ABC`

_sources/autoapi/ocpmodels/trainers/index.rst

@@ -28,7 +28,7 @@ Classes
 
 
 
-.. py:class:: BaseTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier: str, timestamp_id: Optional[str] = None, run_dir: Optional[str] = None, is_debug: bool = False, print_every: int = 100, seed: Optional[int] = None, logger: str = 'tensorboard', local_rank: int = 0, amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm={}, noddp: bool = False)
+.. py:class:: BaseTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier: str, timestamp_id: Optional[str] = None, run_dir: Optional[str] = None, is_debug: bool = False, print_every: int = 100, seed: Optional[int] = None, logger: str = 'wandb', local_rank: int = 0, amp: bool = False, cpu: bool = False, name: str = 'ocp', slurm={}, noddp: bool = False)
 
 
    Bases: :py:obj:`abc.ABC`
@@ -98,7 +98,7 @@ Classes
 
 
 
-.. py:class:: OCPTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='tensorboard', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
+.. py:class:: OCPTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='wandb', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
 
 
    Bases: :py:obj:`ocpmodels.trainers.base_trainer.BaseTrainer`
@@ -141,7 +141,7 @@ Classes
                 (default: :obj:`None`)
    :type seed: int, optional
    :param logger: Type of logger to be used.
-                  (default: :obj:`tensorboard`)
+                  (default: :obj:`wandb`)
    :type logger: str, optional
    :param local_rank: Local rank of the process, only applicable for distributed training.
                       (default: :obj:`0`)

_sources/autoapi/ocpmodels/trainers/ocp_trainer/index.rst

@@ -25,7 +25,7 @@ Classes
 
 
 
-.. py:class:: OCPTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='tensorboard', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
+.. py:class:: OCPTrainer(task, model, outputs, dataset, optimizer, loss_fns, eval_metrics, identifier, timestamp_id=None, run_dir=None, is_debug=False, print_every=100, seed=None, logger='wandb', local_rank=0, amp=False, cpu=False, slurm={}, noddp=False, name='ocp')
 
 
    Bases: :py:obj:`ocpmodels.trainers.base_trainer.BaseTrainer`
@@ -68,7 +68,7 @@ Classes
                 (default: :obj:`None`)
    :type seed: int, optional
    :param logger: Type of logger to be used.
-                  (default: :obj:`tensorboard`)
+                  (default: :obj:`wandb`)
    :type logger: str, optional
    :param local_rank: Local rank of the process, only applicable for distributed training.
                       (default: :obj:`0`)

_sources/core/MODELS.md

@@ -1,28 +1,17 @@
-# Pretrained OCP models
+# Pretrained OCP model checkpoints
 
 This page summarizes all the pretrained models released as part of the [Open Catalyst Project](https://opencatalystproject.org/). All models were trained using this codebase.
 
-
-* [Open Catalyst 2020 (OC20)](#open-catalyst-2020-oc20)
-    * [S2EF models optimized for EFwT](#s2ef-models-optimized-for-efwt)
-    * [S2EF models optimized for force](#s2ef-models-optimized-for-force-only)
-    * [IS2RE models](#is2re-models)
-* [Open Catalyst 2022 (OC22)](#open-catalyst-2022-oc22)
-    * [S2EF total models](#s2ef-total-models)
-* [Open Direct Air Capture 2023 (ODAC23)](#open-direct-air-capture-2023-odac23)
-    * [S2EF models](#s2ef-models)
-    * [IS2RE Direct models](#is2re-direct-models)
-
 * * *
 
 
-# Open Catalyst 2020 (OC20)
+## Open Catalyst 2020 (OC20)
 
 * All configurations for these models are available in the [`configs/`](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs) directory.
 * All of these models are trained on various splits of the OC20 S2EF / IS2RE datasets. For details, see [https://arxiv.org/abs/2010.09990](https://arxiv.org/abs/2010.09990) and https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md.
 * All OC20 models are trained on adsorption energies, i.e. the DFT total energies minus the clean surface and gas phase adsorbate energies. For details on how to train models on OC20 total energies, please read the [referencing section here](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md#oc20-reference-information).
 
-## S2EF models: optimized for EFwT
+### S2EF models: optimized for EFwT
 
 |Model	|Split	|Download	|val ID force MAE (eV / Å)	|val ID EFwT	|
 |---	|---	|---	|---	|---	|
@@ -60,7 +49,7 @@ This page summarizes all the pretrained models released as part of the [Open Cat
 |EquiformerV2 (31M)     |All+MD     |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_31M_ec4_allmd.pt) \| [config](https://github.com/Open-Catalyst-Project/ocp/blob/main/configs/s2ef/all/equiformer_v2/equiformer_v2_N@8_L@4_M@2_31M.yml)  |0.0142  |6.20%  |
 |EquiformerV2 (153M)    |All+MD     |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_153M_ec4_allmd.pt) \| [config](https://github.com/Open-Catalyst-Project/ocp/blob/main/configs/s2ef/all/equiformer_v2/equiformer_v2_N@20_L@6_M@3_153M.yml)    |0.0126     |8.90%  |
 
-## S2EF models: optimized for force only
+### S2EF models: optimized for force only
 
 |Model	|Split	|Download	|val ID force MAE	|
 |---	|---	|---	|---	|
@@ -70,7 +59,7 @@ This page summarizes all the pretrained models released as part of the [Open Cat
 |DimeNet++	|20M+Rattled    |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2021_02/s2ef/dimenetpp_20M_rattled_forceonly.pt)	|0.0614	|
 |DimeNet++	|20M+MD         |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2021_02/s2ef/dimenetpp_20M_md_forceonly.pt)	|0.0594	|
 
-## IS2RE models
+### IS2RE models
 
 |Model	|Split	|Download	|val ID energy MAE	|
 |---	|---	|---	|---	|
@@ -103,14 +92,14 @@ OC20 dataset or pretrained models, as well as the original paper for each model:
 }
 ```
 
-# Open Catalyst 2022 (OC22)
+## Open Catalyst 2022 (OC22)
 
 * All configurations for these models are available in the [`configs/oc22`](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/oc22) directory.
 * All of these models are trained on various splits of the OC22 S2EF / IS2RE datasets. For details, see [https://arxiv.org/abs/2206.08917](https://arxiv.org/abs/2206.08917) and https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md.
 * All OC22 models released here are trained on DFT total energies, in contrast to the OC20 models listed above, which are trained on adsorption energies.
 
 
-## S2EF-Total models
+### S2EF-Total models
 
 |Model	|Training	|Download	|val ID force MAE	|val ID energy MAE	|
 |---	|---	|---	|---	|---	|
@@ -139,7 +128,7 @@ OC22 dataset or pretrained models, as well as the original paper for each model:
 
 * All config files for the ODAC23 models are available in the [`configs/odac`](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/odac) directory.
 
-## S2EF models
+### S2EF models
 
 |Model	|Checkpoint	| Config |
 |---	|---	|---	|
@@ -151,7 +140,7 @@ OC22 dataset or pretrained models, as well as the original paper for each model:
 |EquiformerV2 | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231116/eqv2_31M.pt) | [config](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/odac/s2ef/eqv2_31M.yml) |
 |EquiformerV2 (Large) | [checkpoint](https://dl.fbaipublicfiles.com/dac/checkpoints_20231018/Equiformer_V2_Large.pt) | [config](https://github.com/Open-Catalyst-Project/ocp/tree/main/configs/odac/s2ef/eqv2_153M.yml) |
 
-## IS2RE Direct models
+### IS2RE Direct models
 
 |Model	|Checkpoint	| Config |
 |---	|---	| --- |
@@ -161,7 +150,7 @@ OC22 dataset or pretrained models, as well as the original paper for each model:
 
 The models in the table above were trained to predict relaxed energy directly. Relaxed energies can also be predicted by running structural relaxations using the S2EF models from the previous section.
 
-## IS2RS
+### IS2RS
 
 The IS2RS is solved by running structural relaxations using the S2EF models from the prior section.
 

autoapi/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/data_parallel/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/distutils/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/flags/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/gp_utils/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/hpo_utils/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/logger/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/model_registry/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">

autoapi/ocpmodels/common/registry/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">
@@ -728,8 +726,8 @@ <h3>Attributes<a class="headerlink" href="#attributes" title="Link to this headi
 <p>Usage:</p>
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">ocpmodels.common.registry</span> <span class="kn">import</span> <span class="n">registry</span>
 
-<span class="nd">@registry</span><span class="o">.</span><span class="n">register_logger</span><span class="p">(</span><span class="s2">&quot;tensorboard&quot;</span><span class="p">)</span>
-<span class="k">class</span> <span class="nc">WandB</span><span class="p">():</span>
+<span class="nd">@registry</span><span class="o">.</span><span class="n">register_logger</span><span class="p">(</span><span class="s2">&quot;wandb&quot;</span><span class="p">)</span>
+<span class="k">class</span> <span class="nc">WandBLogger</span><span class="p">():</span>
     <span class="o">...</span>
 </pre></div>
 </div>

autoapi/ocpmodels/common/relaxation/ase_utils/index.html

@@ -188,9 +188,7 @@
 
 <li class="toctree-l1"><a class="reference internal" href="../../../../../core/datasets/oc22.html">Open Catalyst 2022 (OC22)</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../../../core/datasets/odac.html">Open Direct Air Capture 2023 (ODAC23)</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../../../core/MODELS.html">Pretrained OCP models</a></li>
-
-
+<li class="toctree-l1"><a class="reference internal" href="../../../../../core/MODELS.html">Pretrained OCP model checkpoints</a></li>
 </ul>
 <p aria-level="2" class="caption" role="heading"><span class="caption-text">Training your own models</span></p>
 <ul class="nav bd-sidenav">