deactivate autograd during validation (#560)

d7302cc3 · Ignacio Pickering · GitHub · d3847898 · d7302cc3 · d7302cc3
Unverified Commit d7302cc3 authored Dec 16, 2020 by Ignacio Pickering Committed by GitHub Dec 16, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 20 additions and 14 deletions

examples/nnp_training.py examples/nnp_training.py +10 -7

examples/nnp_training_force.py examples/nnp_training_force.py +10 -7

No files found.
--- a/examples/nnp_training.py
+++ b/examples/nnp_training.py
@@ -294,13 +294,16 @@ def validate():
    mse_sum = torch.nn.MSELoss(reduction='sum')
    total_mse = 0.0
    count = 0
-    for properties in validation:
+    model.train(False)
-        species = properties['species'].to(device)
+    with torch.no_grad():
-        coordinates = properties['coordinates'].to(device).float()
+        for properties in validation:
-        true_energies = properties['energies'].to(device).float()
+            species = properties['species'].to(device)
-        _, predicted_energies = model((species, coordinates))
+            coordinates = properties['coordinates'].to(device).float()
-        total_mse += mse_sum(predicted_energies, true_energies).item()
+            true_energies = properties['energies'].to(device).float()
-        count += predicted_energies.shape[0]
+            _, predicted_energies = model((species, coordinates))
+            total_mse += mse_sum(predicted_energies, true_energies).item()
+            count += predicted_energies.shape[0]
+    model.train(True)
    return hartree2kcalmol(math.sqrt(total_mse / count))

--- a/examples/nnp_training_force.py
+++ b/examples/nnp_training_force.py
@@ -208,13 +208,16 @@ def validate():
    mse_sum = torch.nn.MSELoss(reduction='sum')
    total_mse = 0.0
    count = 0
-    for properties in validation:
+    model.train(False)
-        species = properties['species'].to(device)
+    with torch.no_grad():
-        coordinates = properties['coordinates'].to(device).float()
+        for properties in validation:
-        true_energies = properties['energies'].to(device).float()
+            species = properties['species'].to(device)
-        _, predicted_energies = model((species, coordinates))
+            coordinates = properties['coordinates'].to(device).float()
-        total_mse += mse_sum(predicted_energies, true_energies).item()
+            true_energies = properties['energies'].to(device).float()
-        count += predicted_energies.shape[0]
+            _, predicted_energies = model((species, coordinates))
+            total_mse += mse_sum(predicted_energies, true_energies).item()
+            count += predicted_energies.shape[0]
+    model.train(True)
    return hartree2kcalmol(math.sqrt(total_mse / count))