Introduced minor modification in a copy of sobolev, toward debuging it.

sobolev_grad.py

+# A minimal implemetation of sobolev training
+
+import numpy as np
+import torch
+from neural_network import Model
+from datagen import dataGenerator
+import torch.autograd.functional as F
+import matplotlib.pyplot as plt
+
+
+
+# ..............................................................................
+
+
+
+EPOCHS                = 50000                        # Number of Epochs
+lr                    = 1e-3                       # Learning rate
+number_of_batches     = 1                         # Number of batches per epoch
+
+
+
+function_name         = 'simple_bumps'                   # See datagen.py or function_definitions.py for other functions to use
+number_of_data_points = 5
+
+
+
+#.............................................................................
+
+
+
+
+
+X,Y,DY,D2Y            = dataGenerator(function_name, number_of_data_points)
+dataset               = torch.utils.data.TensorDataset(X,Y,DY,D2Y)
+dataloader            = torch.utils.data.DataLoader(dataset, batch_size = number_of_data_points // number_of_batches, 
+                        shuffle=True, num_workers=4)
+
+
+network   = Model(ninput=X.shape[1])
+optimizer = torch.optim.Adam(params = network.parameters(), lr = lr)
+
+
+
+
+epoch_loss_in_value = []
+epoch_loss_in_der1  = []
+epoch_loss_in_der2  = []
+
+
+
+for epoch in range(EPOCHS):
+    network.train()
+    batch_loss_in_value = 0
+    batch_loss_in_der1  = 0
+    batch_loss_in_der2  = 0
+    for idx,(data) in enumerate(dataloader):
+
+        x,y,dy,d2y = data
+        
+        y_hat  = network(x)
+        
+        dy_hat  = torch.vstack( [ F.jacobian(network, state).squeeze() for state in x ] )   # Gradient of net
+        #d2y_hat = torch.stack( [ F.hessian(network, state).squeeze() for state in x ] )     # Hessian of net
+        
+        
+        loss1   = torch.nn.functional.mse_loss(y_hat,y)
+        loss2   = torch.nn.functional.mse_loss(dy_hat, dy)
+        loss3   = 0#torch.nn.functional.mse_loss(d2y_hat, d2y) 
+
+        loss    = loss1 + 10*loss2 + loss3                            # Can add a sobolev factor to give weight to each loss term.
+                                                                   # But it does not really change anything     
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        batch_loss_in_value += loss1.item()
+        batch_loss_in_der1  += loss2.item()
+        #batch_loss_in_der2  += loss3.item()
+
+    epoch_loss_in_value.append( batch_loss_in_value / number_of_batches )
+    epoch_loss_in_der1.append( batch_loss_in_der1 / number_of_batches )
+    #epoch_loss_in_der2.append( batch_loss_in_der2 / number_of_batches )
+
+    
+    if epoch % 10 == 0:
+            print(f"EPOCH : {epoch}")
+            print(f"Loss Values:  {loss1.item()}, Loss Grad : {loss2.item()}") #, Loss Hessian : {loss3.item()}")
+
+plt.ion()
+
+fig, (ax1, ax2, ax3) = plt.subplots(1,3)
+fig.suptitle(function_name.upper())
+
+ax1.semilogy(range(len(epoch_loss_in_value)), epoch_loss_in_value, c = "red")
+#ax2.semilogy(range(len(epoch_loss_in_der1)), epoch_loss_in_der1, c = "green")
+#ax3.semilogy(range(len(epoch_loss_in_der2)), epoch_loss_in_der2, c = "orange")
+
+ax1.set(title='Loss in Value')
+ax2.set(title='Loss in Gradient')
+ax3.set(title='Loss in Hessian')
+
+ax1.set_ylabel('Loss')
+ax1.set_xlabel('Epochs')
+ax2.set_xlabel('Epochs')
+ax3.set_xlabel('Epochs')
+
+
+
+fig.tight_layout()
+#plt.savefig((f"./images/{function_name}.png"))
+#plt.show()
+
+
+#xplt,yplt,dyplt,_            = dataGenerator(function_name, 10000)
+#np.save('plt2.npy',{ "x": xplt.numpy(),"y": yplt.numpy(),"dy": dyplt.numpy()})
+LOAD = np.load( 'plt2.npy',allow_pickle=True).flat[0]
+xplt = torch.tensor(LOAD['x'])
+yplt = torch.tensor(LOAD['y'])
+dyplt = torch.tensor(LOAD['dy'])
+
+ypred = network(xplt)
+
+plt.figure()
+plt.subplot(131)
+plt.scatter(xplt[:,0],xplt[:,1],c=yplt[:,0])
+plt.subplot(132)
+plt.scatter(xplt[:,0],xplt[:,1],c=ypred[:,0].detach())
+plt.subplot(133)
+plt.scatter(xplt[:,0],xplt[:,1],c=(ypred-yplt)[:,0].detach())
+plt.colorbar()

sobolev_training.py

@@ -13,14 +13,14 @@ import matplotlib.pyplot as plt
 
 
 
-EPOCHS                = 500                        # Number of Epochs
+EPOCHS                = 50000                        # Number of Epochs
 lr                    = 1e-3                       # Learning rate
-number_of_batches     = 10                         # Number of batches per epoch
+number_of_batches     = 1                         # Number of batches per epoch
 
 
 
-function_name         = 'ackley'                   # See datagen.py or function_definitions.py for other functions to use
-number_of_data_points = 200
+function_name         = 'simple_bumps'                   # See datagen.py or function_definitions.py for other functions to use
+number_of_data_points = 5
 
 
 
@@ -36,7 +36,7 @@ dataloader            = torch.utils.data.DataLoader(dataset, batch_size = number
                         shuffle=True, num_workers=4)
 
 
-network   = Model()
+network   = Model(ninput=X.shape[1])
 optimizer = torch.optim.Adam(params = network.parameters(), lr = lr)
 
 
@@ -60,14 +60,14 @@ for epoch in range(EPOCHS):
         y_hat  = network(x)
         
         dy_hat  = torch.vstack( [ F.jacobian(network, state).squeeze() for state in x ] )   # Gradient of net
-        d2y_hat = torch.stack( [ F.hessian(network, state).squeeze() for state in x ] )     # Hessian of net
+        #d2y_hat = torch.stack( [ F.hessian(network, state).squeeze() for state in x ] )     # Hessian of net
         
         
         loss1   = torch.nn.functional.mse_loss(y_hat,y)
         loss2   = torch.nn.functional.mse_loss(dy_hat, dy)
-        loss3   = torch.nn.functional.mse_loss(d2y_hat, d2y) 
+        loss3   = 0#torch.nn.functional.mse_loss(d2y_hat, d2y) 
 
-        loss    = loss1 + loss2 + loss3                            # Can add a sobolev factor to give weight to each loss term.
+        loss    = loss1 + 10*loss2 + loss3                            # Can add a sobolev factor to give weight to each loss term.
                                                                    # But it does not really change anything     
         optimizer.zero_grad()
         loss.backward()
@@ -75,24 +75,25 @@ for epoch in range(EPOCHS):
 
         batch_loss_in_value += loss1.item()
         batch_loss_in_der1  += loss2.item()
-        batch_loss_in_der2  += loss3.item()
+        #batch_loss_in_der2  += loss3.item()
 
     epoch_loss_in_value.append( batch_loss_in_value / number_of_batches )
     epoch_loss_in_der1.append( batch_loss_in_der1 / number_of_batches )
-    epoch_loss_in_der2.append( batch_loss_in_der2 / number_of_batches )
+    #epoch_loss_in_der2.append( batch_loss_in_der2 / number_of_batches )
 
     
     if epoch % 10 == 0:
             print(f"EPOCH : {epoch}")
-            print(f"Loss Values:  {loss1.item()}, Loss Grad : {loss2.item()}, Loss Hessian : {loss3.item()}")
+            print(f"Loss Values:  {loss1.item()}, Loss Grad : {loss2.item()}") #, Loss Hessian : {loss3.item()}")
 
+plt.ion()
 
 fig, (ax1, ax2, ax3) = plt.subplots(1,3)
 fig.suptitle(function_name.upper())
 
 ax1.semilogy(range(len(epoch_loss_in_value)), epoch_loss_in_value, c = "red")
-ax2.semilogy(range(len(epoch_loss_in_der1)), epoch_loss_in_der1, c = "green")
-ax3.semilogy(range(len(epoch_loss_in_der2)), epoch_loss_in_der2, c = "orange")
+#ax2.semilogy(range(len(epoch_loss_in_der1)), epoch_loss_in_der1, c = "green")
+#ax3.semilogy(range(len(epoch_loss_in_der2)), epoch_loss_in_der2, c = "orange")
 
 ax1.set(title='Loss in Value')
 ax2.set(title='Loss in Gradient')
@@ -106,5 +107,24 @@ ax3.set_xlabel('Epochs')
 
 
 fig.tight_layout()
-plt.savefig((f"./images/{function_name}.png"))
-plt.show()
+#plt.savefig((f"./images/{function_name}.png"))
+#plt.show()
+
+
+#xplt,yplt,dyplt,_            = dataGenerator(function_name, 10000)
+#np.save('plt2.npy',{ "x": xplt.numpy(),"y": yplt.numpy(),"dy": dyplt.numpy()})
+LOAD = np.load( 'plt2.npy',allow_pickle=True).flat[0]
+xplt = torch.tensor(LOAD['x'])
+yplt = torch.tensor(LOAD['y'])
+dyplt = torch.tensor(LOAD['dy'])
+
+ypred = network(xplt)
+
+plt.figure()
+plt.subplot(131)
+plt.scatter(xplt[:,0],xplt[:,1],c=yplt[:,0])
+plt.subplot(132)
+plt.scatter(xplt[:,0],xplt[:,1],c=ypred[:,0].detach())
+plt.subplot(133)
+plt.scatter(xplt[:,0],xplt[:,1],c=(ypred-yplt)[:,0].detach())
+plt.colorbar()