I am using google colab to implement a Physics-Informed Neural Network. I upgraded my Colab to System RAM of 12.7 GB and Disk of 107.7 GB. However, when computing the loss function, I am getting the following error: Your session crashed after using all available RAM.
This is my code:
`
def compute_loss(model, grid_points, real_OSM_values_Circle_tensor):
#Compute physics based loss
physics_based_loss = tf.reduce_mean(tf.square(get_residual(model,grid_points)))
#initialise loss
loss = physics_based_loss
#Add physics_based_loss and data loss
for i in range(len(real_OSM_values_Circle_tensor)):
u_pred = model(grid_points[i:i+1,:])
loss += tf.reduce_mean(tf.square(u_pred - real_OSM_values_Circle_tensor[i]))
return loss
#model = init_model()
#compute_loss(model, grid_points, real_OSM_values_Circle_tensor)
def get_grad(model,grid_points):
with tf.GradientTape(persistent=True) as tape:
#This tape is for derivatives with respect to trainable variables
tape.watch(model.trainable_variables)
loss = compute_loss(model, grid_points, real_OSM_values_Circle_tensor)
g = tape.gradient(loss, model.trainable_variables)
del tape
return loss, g
lr = 0.001
Choose the optimizer
optim = tf.keras.optimizers.Adam(learning_rate=lr)
Define one training step as a TensorFlow function to increase speed of training
@tf.function
def train_step():
Compute current loss and gradient w.r.t. parameters
loss, grad_theta = get_grad(model,grid_points)
Perform gradient descent step
optim.apply_gradients(zip(grad_theta, model.trainable_variables))
return loss
Number of training epochs
N = 1
hist = []
Start timer
t0 = time()
for i in range(N+1):
loss = train_step()
Append current loss to hist
hist.append(loss.numpy())
“”"
Output current loss after 50 iterates
if i%50 == 0:
print(‘It {:05d}: loss = {:10.8e}’.format(i,loss))
“”"
print(‘It {:05d}: loss = {:10.8e}’.format(i,loss))
Print computation time
print(‘\nComputation time: {} seconds’.format(time()-t0))
`
When I try initialising a model and running it, the code seems to get stuck on line loss = train_step(). I tried setting up a really simple NN and the number of epochs to 1 just to ensure the code runs but even then I am getting the same error. Should I upgrade the RAM or is there a way to make my code more efficient (len(real_OSM_values_Circle_tensor) = 22500).
This is my first time in this forum, so please let me know if this is not the appropriate place and where can I post this question. Thank you!