diff --git a/examples/ogbn_train.py b/examples/ogbn_train.py index 56a5c1c7a538..2f02adac5f13 100644 --- a/examples/ogbn_train.py +++ b/examples/ogbn_train.py @@ -33,7 +33,7 @@ action='store_true', help='Whether or not to use GAT model', ) -parser.add_argument('-e', '--epochs', type=int, default=10) +parser.add_argument('-e', '--epochs', type=int, default=50) parser.add_argument('--num_layers', type=int, default=3) parser.add_argument('--num_heads', type=int, default=2, help='number of heads for GAT model.') @@ -179,6 +179,7 @@ def test(loader: NeighborLoader) -> float: lr=args.lr, weight_decay=args.wd, ) +scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max',patience=5) print(f'Total time before training begins took ' f'{time.perf_counter() - wall_clock_start:.4f}s') @@ -204,6 +205,10 @@ def test(loader: NeighborLoader) -> float: if val_acc > best_val: best_val = val_acc times.append(time.perf_counter() - train_start) + for param_group in optimizer.param_groups: + print('lr:') + print(param_group['lr']) + scheduler.step(val_acc) print(f'Average Epoch Time on training: ' f'{torch.tensor(train_times).mean():.4f}s')