Commit 124ce875 by lvzhengyang

add test

parent 2717b7cc
......@@ -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)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment