add test

build_model/load_graph.py

@@ -20,18 +20,13 @@ dir_prefix = f"../data/{pdk}/{tag}"
     
 # blocks = "aes aes-mbff ethmac gcd ibex jpeg mock-alu uart".split()
 blocks = "aes aes-mbff gcd ibex jpeg uart".split()
-# blocks_test = set(('aes', 'ibex'))
-blocks_test = set()
+blocks_test = set('uart'.split())
+tag = '6'
 blocks_train = set()
 for block in blocks:
     if not block in blocks_test:
         blocks_train.add(block)
 
-# train config
-netdelay = True
-celldelay = True
-groundtruth = True
-
 def gen_topo(g_hetero):
     torch.cuda.synchronize()
     time_s = time.time()
@@ -71,11 +66,11 @@ def load_data():
         data[block] = g, ts
     return data
 
-def train(model, data_train):
+def train(model, data_train, data_test):
     writer = SummaryWriter()
     optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
 
-    batch_size = 5
+    batch_size = min(5, len(data_train))
     for e in range(100000):
         model.train()
         train_loss_tot_net_delays, train_loss_tot_cell_delays, train_loss_tot_ats = 0, 0, 0
@@ -106,14 +101,28 @@ def train(model, data_train):
         train_loss_epoch_net_delays, train_loss_epoch_cell_delays = 0, 0
         
         # log
-        if (e + 1) % 10 == 0:
-            print('epoch: {}, net_delay_loss (train): {:.6f}, cell_delay_loss (train): {:.6f}'.format(
-                e + 1, train_loss_tot_net_delays, train_loss_tot_cell_delays
+        if (e + 1) % 20 == 0:
+            # test model
+            with torch.no_grad():
+                model.eval()
+                test_loss_tot_net_delays, test_loss_tot_cell_delays, test_loss_tot_ats = 0, 0, 0
+                for k, (g, ts) in data_test.items():
+                    pred_net_delays, pred_cell_delays, pred_atslew = model(g, ts)
+                    test_loss_tot_net_delays += F.mse_loss(pred_net_delays, g.ndata['n_net_delays_log']).item()
+                    test_loss_tot_cell_delays += F.mse_loss(pred_cell_delays, g.edges['cell_out'].data['e_cell_delays_log']).item()
+                print('epoch: {}, net_delay_loss (train): {:.6e}, cell_delay_loss (train): {:.6e}, net_delay_loss (test): {:.6e}, cell_delay_loss (test): {:.6e}'.format(
+                    e + 1, train_loss_tot_net_delays, train_loss_tot_cell_delays, test_loss_tot_net_delays, test_loss_tot_cell_delays
                 ))
+                writer.add_scalar('net_delays_loss/test: ', test_loss_tot_net_delays, e)
+                writer.add_scalar('cell_delays_loss/test: ', test_loss_tot_cell_delays, e)
+        
         # save model
-        if (e + 1) % 100 == 0:
+        if (e + 1) % 200 == 0:
             print('-------- Save Model --------')
-            save_path = os.path.join('weights', f'{e}.pt')
+            save_dir = os.path.join('weights', tag)
+            save_path = os.path.join(save_dir, '%d-%.3f-%.3f-%.3f-%.3f.pt'%(
+                e + 1, train_loss_tot_net_delays, train_loss_tot_cell_delays, test_loss_tot_net_delays, test_loss_tot_cell_delays))
+            os.makedirs(save_dir, exist_ok=True)
             torch.save(model.state_dict(), save_path)
     writer.close()
 
@@ -125,5 +134,5 @@ if __name__ == "__main__":
     model = PredModel()
     model.cuda()
 
-    train(model, data_train)
+    train(model, data_train, data_test)