Energy forces

hydragnn/models/Base.py

@@ -15,6 +15,7 @@
 from torch_geometric.nn import global_mean_pool, BatchNorm
 from torch.nn import GaussianNLLLoss
 from torch.utils.checkpoint import checkpoint
+import torch_scatter
 from hydragnn.utils.model import activation_function_selection, loss_function_selection
 import sys
 from hydragnn.utils.distributed import get_device
@@ -355,6 +356,60 @@ def loss(self, pred, value, head_index):
         elif self.ilossweights_hyperp == 1:
             return self.loss_hpweighted(pred, value, head_index, var=var)
 
+    def energy_force_loss(self, pred, data):
+        # Asserts
+        assert (
+            data.pos is not None and data.energy is not None and data.forces is not None
+        ), "data.pos, data.energy, data.forces must be provided for energy-force loss. Check your dataset creation and naming."
+        assert (
+            data.pos.requires_grad
+        ), "data.pos does not have grad, so force predictions cannot be computed. Check that data.pos has grad set to true before prediction."
+        assert (
+            self.num_heads == 1 and self.head_type[0] == "node"
+        ), "Force predictions are only supported for models with one head that predict nodal energy. Check your num_heads and head_types."
+        # Initialize loss
+        tot_loss = 0
+        tasks_loss = []
+        # Energies
+        node_energy_pred = pred[0]
+        graph_energy_pred = torch_scatter.scatter_add(
+            node_energy_pred, data.batch, dim=0
+        ).float()
+        graph_energy_true = data.energy
+        energy_loss_weight = self.loss_weights[
+            0
+        ]  # There should only be one loss-weight for energy
+        tot_loss += (
+            self.loss_function(graph_energy_pred, graph_energy_true)
+            * energy_loss_weight
+        )
+        tasks_loss.append(self.loss_function(graph_energy_pred, graph_energy_true))
+        # Forces
+        forces_true = data.forces.float()
+        forces_pred = torch.autograd.grad(
+            graph_energy_pred,
+            data.pos,
+            grad_outputs=torch.ones_like(graph_energy_pred),
+            retain_graph=graph_energy_pred.requires_grad,  # Retain graph only if needed (it will be needed during training, but not during validation/testing)
+            create_graph=True,
+        )[0].float()
+        assert (
+            forces_pred is not None
+        ), "No gradients were found for data.pos. Does your model use positions for prediction?"
+        forces_pred = -forces_pred
+        force_loss_weight = (
+            energy_loss_weight
+            * torch.mean(torch.abs(graph_energy_true))
+            / (torch.mean(torch.abs(forces_true)) + 1e-8)
+        )  # Weight force loss and graph energy equally
+        tot_loss += (
+            self.loss_function(forces_pred, forces_true) * force_loss_weight
+        )  # Have force-weight be the complement to energy-weight
+        ## FixMe: current loss functions require the number of heads to be the number of things being predicted
+        ##        so, we need to do loss calculation manually without calling the other functions.
+
+        return tot_loss, tasks_loss
+
     def loss_nll(self, pred, value, head_index, var=None):
         # negative log likelihood loss
         # uncertainty to weigh losses in https://openaccess.thecvf.com/content_cvpr_2018/papers/Kendall_Multi-Task_Learning_Using_CVPR_2018_paper.pdf

hydragnn/run_training.py

@@ -175,6 +175,7 @@ def _(config: dict, use_deepspeed=False):
         plot_hist_solution,
         create_plots,
         use_deepspeed=use_deepspeed,
+        compute_grad_energy=config["NeuralNetwork"]["Training"]["compute_grad_energy"],
     )
 
     save_model(model, optimizer, log_name, use_deepspeed=use_deepspeed)

hydragnn/train/train_validate_test.py

@@ -66,6 +66,7 @@ def train_validate_test(
     plot_hist_solution=False,
     create_plots=False,
     use_deepspeed=False,
+    compute_grad_energy=False,
 ):
     num_epoch = config["Training"]["num_epoch"]
     EarlyStop = (
@@ -162,6 +163,7 @@ def train_validate_test(
                 verbosity,
                 profiler=prof,
                 use_deepspeed=use_deepspeed,
+                compute_grad_energy=compute_grad_energy,
             )
             tr.stop("train")
             tr.disable()
@@ -172,14 +174,19 @@ def train_validate_test(
             continue
 
         val_loss, val_taskserr = validate(
-            val_loader, model, verbosity, reduce_ranks=True
+            val_loader,
+            model,
+            verbosity,
+            reduce_ranks=True,
+            compute_grad_energy=compute_grad_energy,
         )
         test_loss, test_taskserr, true_values, predicted_values = test(
             test_loader,
             model,
             verbosity,
             reduce_ranks=True,
             return_samples=plot_hist_solution,
+            compute_grad_energy=compute_grad_energy,
         )
         scheduler.step(val_loss)
         if writer is not None:
@@ -434,7 +441,15 @@ def gather_tensor_ranks(head_values):
     return head_values
 
 
-def train(loader, model, opt, verbosity, profiler=None, use_deepspeed=False):
+def train(
+    loader,
+    model,
+    opt,
+    verbosity,
+    profiler=None,
+    use_deepspeed=False,
+    compute_grad_energy=False,
+):
     if profiler is None:
         profiler = Profiler()
 
@@ -492,8 +507,13 @@ def train(loader, model, opt, verbosity, profiler=None, use_deepspeed=False):
             data = data.to(get_device())
             if trace_level > 0:
                 tr.stop("h2d", **syncopt)
-            pred = model(data)
-            loss, tasks_loss = model.module.loss(pred, data.y, head_index)
+            if compute_grad_energy:  # for force and energy prediction
+                data.pos.requires_grad = True
+                pred = model(data)
+                loss, tasks_loss = model.module.energy_force_loss(pred, data)
+            else:
+                pred = model(data)
+                loss, tasks_loss = model.module.loss(pred, data.y, head_index)
             if trace_level > 0:
                 tr.start("forward_sync", **syncopt)
                 MPI.COMM_WORLD.Barrier()
@@ -541,7 +561,7 @@ def train(loader, model, opt, verbosity, profiler=None, use_deepspeed=False):
 
 
 @torch.no_grad()
-def validate(loader, model, verbosity, reduce_ranks=True):
+def validate(loader, model, verbosity, reduce_ranks=True, compute_grad_energy=False):
 
     total_error = torch.tensor(0.0, device=get_device())
     tasks_error = torch.zeros(model.module.num_heads, device=get_device())
@@ -565,8 +585,14 @@ def validate(loader, model, verbosity, reduce_ranks=True):
             loader.dataset.ddstore.epoch_end()
         head_index = get_head_indices(model, data)
         data = data.to(get_device())
-        pred = model(data)
-        error, tasks_loss = model.module.loss(pred, data.y, head_index)
+        if compute_grad_energy:  # for force and energy prediction
+            with torch.enable_grad():
+                data.pos.requires_grad = True
+                pred = model(data)
+                error, tasks_loss = model.module.energy_force_loss(pred, data)
+        else:
+            pred = model(data)
+            error, tasks_loss = model.module.loss(pred, data.y, head_index)
         total_error += error * data.num_graphs
         num_samples_local += data.num_graphs
         for itask in range(len(tasks_loss)):
@@ -585,7 +611,14 @@ def validate(loader, model, verbosity, reduce_ranks=True):
 
 
 @torch.no_grad()
-def test(loader, model, verbosity, reduce_ranks=True, return_samples=True):
+def test(
+    loader,
+    model,
+    verbosity,
+    reduce_ranks=True,
+    return_samples=True,
+    compute_grad_energy=False,
+):
 
     total_error = torch.tensor(0.0, device=get_device())
     tasks_error = torch.zeros(model.module.num_heads, device=get_device())
@@ -612,8 +645,14 @@ def test(loader, model, verbosity, reduce_ranks=True, return_samples=True):
             loader.dataset.ddstore.epoch_end()
         head_index = get_head_indices(model, data)
         data = data.to(get_device())
-        pred = model(data)
-        error, tasks_loss = model.module.loss(pred, data.y, head_index)
+        if compute_grad_energy:  # for force and energy prediction
+            with torch.enable_grad():
+                data.pos.requires_grad = True
+                pred = model(data)
+                error, tasks_loss = model.module.energy_force_loss(pred, data)
+        else:
+            pred = model(data)
+            error, tasks_loss = model.module.loss(pred, data.y, head_index)
         ## FIXME: temporary
         if int(os.getenv("HYDRAGNN_DUMP_TESTDATA", "0")) == 1:
             if model.module.var_output: