Function train

net_delay_prediction/train.py:63–108

(model, args)

        test_dict(data_test)

def train(model, args):
    data_train, data_test = load_data(args)
    
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    for e in range(args.iteration):
        model.train()
        train_loss_tot_net_delays = 0
        optimizer.zero_grad()

        for k, g in random.sample(data_train.items(), args.batch_size):
            pred_net_delays= model(g)
            loss_net_delays = 0

            loss_net_delays = F.mse_loss(pred_net_delays, g.edges['net_out'].data['net_delays_log'])
            train_loss_tot_net_delays += loss_net_delays.item()
            loss_net_delays.backward()
            
        optimizer.step()

        if e == 0 or e % 20 == 19:
            with torch.no_grad():
                model.eval()
                test_loss_tot_net_delays= 0
                for k, g in data_test.items():
                    pred_net_delays= model(g)

                    test_loss_tot_net_delays += F.mse_loss(pred_net_delays, g.edges['net_out'].data['net_delays_log']).item()

                print('Epoch {}, net delay {:.6f}/{:.6f})'.format(
                    e,
                    train_loss_tot_net_delays / args.batch_size,
                    test_loss_tot_net_delays / len(data_test)
                    )
                    )

            if e == 0 or e % 200 == 199 or (e > 6000 and test_loss_tot_net_delays / len(data_test) < 6):
                if args.checkpoint:
                    save_path = './checkpoints/{}/{}.pth'.format(args.checkpoint, e)
                    torch.save(model.state_dict(), save_path)
                    print('saved model to', save_path)
                try:
                    test_netdelay(model)
                except ValueError as e:
                    print(e)
                    print('Error testing, but ignored')

if __name__ == '__main__':
    args = parser.parse_args()