Loading DPs as TF models and performing external backprop

[siddarthachar]

I wanted to ask about a technical aspect of DeePMD. Is there a way by which I can load DPs as TF models and perform backprop externally, like on a jupyter notebook. I am able to use DPs in Jupyter notebooks for other things but this. I have a project in mind that makes uses of loss function gradients w.r.t atomic and for that I would want to perform backprop along the DP. Using DPs to perform these tasks will be the best for us since we’ve had great success with DPs so far. This is a toy code that I wrote (not the actual project):

import numpy as np
from deepmd import (
    DeepPotential,
 
dp=DeepPotential('nh3.pb')
coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
cell = np.diag(10 * np.ones(3)).reshape([1, -1])
atype = [1,0,1]
e, f, v = dp.eval(coord, cell, atype)
 
tf.config.run_functions_eagerly(True)
 
# Initialize input and output tensors
inputs = tf.constant([[0.2, 0.3, 0.4]])
coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
cell = np.diag(10 * np.ones(3)).reshape([1, -1])
atype = [1,0,1]
targets = tf.constant([[19.0]])
 
# Initialize weights and biases as tensors
weights = tf.Variable(tf.random.normal([3, 1]))
biases = tf.Variable(tf.zeros([1]))
 
# Define the forward pass
def forward_pass(inputs):
    return tf.matmul(inputs, weights) + biases
 
# Define the loss function
def loss_function(predictions, targets):
    e, f, v = dp.eval(coord, cell, atype)
#     print(e)
    return tf.reduce_mean(tf.square(e + predictions - targets))
 
# Define the optimizer
optimizer = tf.optimizers.SGD(learning_rate=0.1)
 
# Perform backpropagation
# inputs = [coord, cell, atype]
for epoch in range(10):
    with tf.GradientTape() as tape:
        # Forward pass
        predictions = forward_pass(inputs)
#         print(predictions, targets)
 
        # Compute the loss
        loss = loss_function(predictions, targets)
    # Compute the gradients
    gradients = tape.gradient(loss, [weights, biases])
    # Update the weights and biases
    optimizer.apply_gradients(zip(gradients, [weights, biases]))
 
    print(f"Epoch {epoch+1}: Loss = {loss}")
 
# Make predictions on new data
new_inputs = tf.constant([[0.1, 0.2, 0.3]])
 
new_predictions = forward_pass(inputs)
print("Predictions:")
print(new_predictions)

This is something that I wrote, but the back prop does not happen. I am assuming that the “dp” does not allow for gradient calculations via tf. Do you know what should be done? I am trying to write an application that uses DPs to perform adversarial attack based active learning. It would be great if there’s a way to backprop through the model in a simple way, if you knew one.

[wanghan-iapcm]

The implementation of deepmd-kit is based on tensorflow v1, with which the eager execution is not supported.
You may get the computational graphs of energy, force and virial by accessing the t_energy, t_force and t_virial members of your DeepPotential object.

[siddarthachar]

Thanks Han! I make changes to my code to maintain tf v1 thourghout and it seems to be working fine now.

# import tensorflow as tf
# import numpy as np
# from deepmd import (
#     DeepPotential,)
tf.compat.v1.disable_eager_execution()
tf.compat.v1.reset_default_graph()

coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
cell = np.diag(10 * np.ones(3)).reshape([1, -1])
atype = [1,0,1]
e, f, v = dp.eval(coord, cell, atype)

# tf.config.run_functions_eagerly(True)

# Initialize input and output tensors
inputs = tf.constant([[0.2, 0.3, 0.4]])
coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
cell = np.diag(10 * np.ones(3)).reshape([1, -1])
atype = [1,0,1]
targets = tf.constant([[19.0]])

# Initialize weights and biases as tensors
weights = tf.Variable(tf.random.normal([3, 1]))
biases = tf.Variable(tf.zeros([1]))

# Define the forward pass
def forward_pass(inputs):
    return tf.matmul(inputs, weights) + biases

# Define the loss function
def loss_function(predictions, targets):
    e, f, v = dp.eval(coord, cell, atype)
#     print(e)
    return tf.reduce_mean(tf.square(e + predictions - targets))

# Define the optimizer
# optimizer = tf.optimizers.SGD(learning_rate=0.1)
optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)

# Perform backpropagation
# inputs = [coord, cell, atype]
for epoch in range(10):
    with tf.GradientTape() as tape:
        # Forward pass
        predictions = forward_pass(inputs)
#         print(predictions, targets)

        # Compute the loss
        loss = loss_function(predictions, targets)
    # Compute the gradients
    gradients = tape.gradient(loss, [weights, biases])
    # Update the weights and biases
    optimizer.apply_gradients(zip(gradients, [weights, biases]))

#     print(f"Epoch {epoch+1}: Loss = {loss}")
    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        print(f"Epoch {epoch+1}: Loss = {sess.run(loss)}")

# Make predictions on new data
new_inputs = tf.constant([[0.1, 0.2, 0.3]])

new_predictions = forward_pass(inputs)
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    print("Predictions:")
    print(sess.run(new_predictions))

It also seems like the loss function is updating.

---

Regarding using dp.t_energy, I can't figure out how to use it to get energy tensors and use it to perform further calculations. Say I want to define a loss function loss = energy**2 + sum(forces) and use it to perform backprop. And I want the values of energy and forces to be calculated using the DeepPotential (dp) object. I could not find a way to get the tensor form of energy and forces without performing dp.eval. How can I use dp.t_energy and dp.t_forces to achieve this? It is of type: <tf.Tensor 'load/o_energy:0' shape=(?,) dtype=float64>, which I think is just a general tensor representation of energy. Ideally, something like loss = dp.t_energy**2 + sum(dp.t_forces) is what I wanted. Do you know what I can do for this? I want to compute d(loss)/dx, where x is the atomic coordinates. This would mean x also needs to be a tensor with a computational graph.

Thanks!