[Python] GNN Loss NaN after first training example?

I am trying to train a GNN but am getting a NaN loss function immediately after the first training example. Below I have included all of the pertinent code. My input is 385 points in 3D space confined to a box with periodic boundary conditions. My model takes 5 inputs:

• a 385 length vector of node labels (either 1, 2, or 3)
• a 385x5x300 tensor: each of the 385 nodes has 5 edge connections made of a 300 component feature vector
• a 385x5 matrix of node indices giving the index of connections to each node
• a 384x3 matrix that represents the 3D points for node types 1 and 2
• a 1x3 vector that represents the 3D points for a specific single node of type 3

I only want to make a prediction at the given node of type 3, which should predict an update to the 3D position of the last input. However, when I try to train the model I get NaN loss after the first training example. I have evaluated the untrained model on the first 5 inputs and calculated the loss of these prediction and they are finite. I would prefer not to do any scaling/normalization of my data, particularly since my inputs are bound to a finite 3D space and my loss function is bound from 0 to half the box length due to the minimum image convention.

My inputs i1-i5 and target ys_list are defined as:

tr_dat_list = [make_graph_in_only_b(i[0:-1], 15.66) for i in pos_g[0:1500]]
ys_list = pos_g[0:1500,-1].reshape((len(pos_g[0:1500]),1,3))

i1 = np.array([i[0] for i in tr_dat_list]).reshape((len(tr_dat_list), 385))
i2 = np.array([i[1] for i in tr_dat_list]).reshape((len(tr_dat_list), 385, 5, 300))
i3 = np.array([i[2] for i in tr_dat_list]).reshape((len(tr_dat_list), 385, 5))
i4 = np.array([i[3] for i in tr_dat_list]).reshape((len(tr_dat_list), 384, 3))
i5 = np.array([i[4] for i in tr_dat_list]).reshape((len(tr_dat_list), 3))


Where pos_g is a list of 386x3 matrices. The first 384 rows are the 3D coordinates of nodes type 1 and 2, while the 385th row is an initial guess for my label (the 386th row).

You can generate some synthetic data to reproduce the error as follows:

pos_g = []
for i in range(1500):
fpos = np.random.rand(384,3)
lp = np.array(list(map(np.mean, fpos[[12, 5, 1]].T)))
lpg = lp + ((np.random.rand(3))-0.5)/5
pos_g.append(np.concatenate((fpos, lpg.reshape((1,3)), lp.reshape((1,3))), axis=0))
pos_g = np.array(pos_g)

tr_dat_list = [make_graph_in_only_b(i[0:-1], 15.66) for i in pos_g[0:1500]]
ys_list = pos_g[0:1500,-1].reshape((len(pos_g[0:1500]),1,3))

i1 = np.array([i[0] for i in tr_dat_list]).reshape((len(tr_dat_list), 385))
i2 = np.array([i[1] for i in tr_dat_list]).reshape((len(tr_dat_list), 385, 5, 300))
i3 = np.array([i[2] for i in tr_dat_list]).reshape((len(tr_dat_list), 385, 5))
i4 = np.array([i[3] for i in tr_dat_list]).reshape((len(tr_dat_list), 384, 3))
i5 = np.array([i[4] for i in tr_dat_list]).reshape((len(tr_dat_list), 3))


Here is the model class:

class net(tf.keras.Model):

def __init__(self, gnn_blocks, channels, label_dim, k_its, boxl, **kwargs):
super(net, self).__init__(**kwargs)
self.gnn_blocks = gnn_blocks
self.nchannels = channels
self.nk = k_its
self.bl = boxl

self.embedding = tf.keras.layers.Embedding(4, channels)
self.h1s = [make_h1(channels) for _ in range(self.gnn_blocks)]
self.h2s = [make_h2(channels) for _ in range(self.gnn_blocks)]
self.h3s = [make_h3(channels) for _ in range(self.gnn_blocks)]
self.readout_l1 = tf.keras.layers.Dense(channels // 2, activation=ssp)
self.readout_l2 = tf.keras.layers.Dense(label_dim)

def build(self, input_shape):

dummy_nodes = tf.zeros(shape=( 385,), dtype=tf.int32)
dummy_edges = tf.zeros(shape=( 385, 5, 300), dtype=tf.float32)
dummy_edge_i = tf.zeros(shape=( 385, 5), dtype=tf.int32)
dummy_wt_cds = tf.zeros(shape=( 384,3), dtype=tf.float32)
dummy_wn_guess = tf.zeros(shape=( 3), dtype=tf.float32)

def call(self, inputs):

guessed_wns = tf.TensorArray(dtype=tf.float32, size=self.nk)

nodes = inputs[0]
edge_features = inputs[1]
edge_i = inputs[2]
wt_cds = inputs[3]
wn_guess_init = inputs[4]

nodes = self.embedding(nodes)
for i in range(self.gnn_blocks):

v_sk = tf.gather(nodes, edge_i, batch_dims=1)
e_k = self.h1s(v_sk) * self.h2s(edge_features)
e_i = tf.reduce_sum(e_k, axis=2)
nodes += self.h3s(e_i)
tf.debugging.check_numerics(v_sk, message='Checking v_sk')
tf.debugging.check_numerics(e_k, message='Checking e_k')
tf.debugging.check_numerics(e_i, message='Checking e_i')
tf.debugging.check_numerics(nodes, message='Checking nodes')

rn = tf.reshape(nodes[:, -1, :], shape=(1,self.nchannels))
rn = self.readout_l1(rn)
rn = self.readout_l2(rn)
upd_wn = mivec((wn_guess_init + rn), self.bl)

k=0
guessed_wns = guessed_wns.write(k, tf.squeeze(upd_wn))
k += 1
upd_wn = tf.reshape(upd_wn, shape=(1,1,3))
new_cds = tf.concat([wt_cds, upd_wn], axis=1)
new_cds = tf.reshape(new_cds, shape=(385,3))

while k < self.nk:
nodes, edge_features, edge_i, wat_cds, wan_guess = make_graph_in_only_b(new_cds, self.bl)
nodes = tf.expand_dims(nodes, 0)
edge_features = tf.expand_dims(edge_features, 0)
edge_i = tf.expand_dims(edge_i, 0)

nodes = self.embedding(nodes)
for i in range(self.gnn_blocks):

v_sk = tf.gather(nodes, edge_i, batch_dims=1)
e_k = self.h1s(v_sk) * self.h2s(edge_features)
e_i = tf.reduce_sum(e_k, axis=2)
nodes += self.h3s(e_i)

rn = tf.reshape(nodes[:, -1, :], shape=(1,self.nchannels))
rn = self.readout_l1(rn)
rn = self.readout_l2(rn)
upd_wn = mivec(wn_guess + rn, self.bl)
guessed_wns = guessed_wans.write(k, tf.squeeze(upd_wn))
upd_wn = tf.reshape(upd_wn, shape=(1,1,3))
wt_cds = tf.reshape(wt_cds, shape=(1,384,3))
new_cds = tf.concat([wt_cds, upd_wn], axis=1)
new_cds = tf.reshape(new_cds, shape=(385,3))
k += 1

result = guessed_wns.stack()
result = tf.reshape(tf.convert_to_tensor(result), shape=(1,k,3))
tf.debugging.check_numerics(result, "NaNs detected in final result")
return result


And here are the functions used in the above model:

def ssp(x):
return tf.math.log(0.5 * tf.math.exp(x) + 0.5)


def make_h1(units):
return tf.keras.Sequential([tf.keras.layers.Dense(units)])


def make_h2(units):
return tf.keras.Sequential(
[
tf.keras.layers.Dense(units, activation=ssp),
tf.keras.layers.Dense(units, activation=ssp),
]
)

def make_h3(units):
return tf.keras.Sequential(
[tf.keras.layers.Dense(units, activation=ssp), tf.keras.layers.Dense(units)]
)

def mivec(v, bl):
return v - bl*tf.math.round(v/bl)

def get_edges(positions, NN, bl, sorted=True):
M = tf.shape(input=positions)[0]
NN = tf.minimum(NN, M)
qexpand = tf.expand_dims(positions, 1) # one column
qTexpand = tf.expand_dims(positions, 0) # one row
qtile = tf.tile(qexpand, [1, M, 1])
qTtile = tf.tile(qTexpand, [M, 1, 1])
dist_mat = qTtile - qtile
dist_mat -= bl*tf.math.round(dist_mat/bl)
dist = tf.norm(tensor=dist_mat, axis=2)
mask = dist >= 5e-4
mask_cast = tf.cast(mask, dtype=dist.dtype)
dist_mat_r = dist * mask_cast + (1 - mask_cast) * 1000
topk = tf.math.top_k(-dist_mat_r, k=NN, sorted=sorted)
return -topk.values, topk.indices

def rbf(r):
return tf.exp(-gamma * (r[..., tf.newaxis] - mu) ** 2)

def make_graph_in_only_b(x, bl=15.66):
wat_node_tens = tf.constant([[1,2,2]])
edge_r, edge_i = get_edges(x, MAX_DEGREE, bl)
edge_features = rbf(edge_r)
elec_node = tf.constant([3])
nodes = tf.concat([tf.reshape(tf.repeat(wat_node_tens, repeats=tf.cast((tf.shape(x)[0]-1)/3, tf.int32), axis=0), shape=(tf.shape(x)[0]-1,)), tf.constant([3])], -1)
return (nodes, edge_features, edge_i, tf.cast(x[0:-1], tf.float32), tf.cast(x[-1], tf.float32))



I am trying to train the model using fit():

ss = net(6, 64, 3, 4, 15.66)

ss.compile(loss=total_loss(1.05,4,15.66),
optimizer=tf.keras.optimizers.Adam(learning_rate=2.0e-3,clipvalue=1.0)
)

result = ss.fit([i1, i2, i3, i4, i5], ys_list, batch_size=1, epochs=10, steps_per_epoch=len(i1), callbacks=[print_on_end()])


Where the inputs [i1, i2, i3, i4, i5] are the outputs of make_graph_in_only_b() operating on a 3D point cloud. With the loss function and callback:

class print_on_end(tf.keras.callbacks.Callback):
def on_batch_end(self, batch, logs={}):
print()

def total_loss(gamma, k, bl):
def iterations_loss(y_true, y_pred):
tf.debugging.check_numerics(y_true, message='Checking y_true')
#y_pred = y_pred[0]
t_loss = 0.0
for i in range(k):
t_dv = y_true - y_pred
t_dv = t_dv - bl*tf.math.round(t_dv/bl)
t_loss += (gamma**(i+1))*tf.norm(t_dv)
tf.debugging.check_numerics(t_loss, message='Checking t_loss')
return tf.cast(t_loss, tf.float32)
return iterations_loss


The first epoch returns a finite loss and the second gives nan. I have successfully predicted outputs on my first few training examples with:

ss = net(6, 64, 3, 4, 15.66)

out_test = ss.predict([i1[0:5], i2[0:5], i3[0:5], i4[0:5], i5[0:5]], batch_size=1)


Thank you in advance for any guidance!

Continue reading...