Update based on comments

Author: yanbing-j

benchmark/citation/train_eval.py

@@ -39,6 +39,9 @@ def random_planetoid_splits(data, num_classes):
 def run_train(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
               profiling, use_compile, permute_masks=None, logger=None):
     val_losses, accs, durations = [], [], []
+    if use_compile:
+        model = torch_geometric.compile(model)
+
     for run in range(runs):
         data = dataset[0]
         if permute_masks is not None:
@@ -97,14 +100,6 @@ def run_train(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
         with torch_profile():
             train(model, optimizer, data)
 
-    if use_compile:
-        print("Using torch.compile")
-        compiled_model = torch_geometric.compile(model)
-        train(compiled_model, optimizer, data)
-        train(compiled_model, optimizer, data)
-        with timeit():
-            train(compiled_model, optimizer, data)
-
 
 @torch.no_grad()
 def run_inference(dataset, model, epochs, profiling, bf16, use_compile,
@@ -115,6 +110,8 @@ def run_inference(dataset, model, epochs, profiling, bf16, use_compile,
     data = data.to(device)
 
     model.to(device).reset_parameters()
+    if use_compile:
+        model = torch_geometric.compile(model)
 
     if torch.cuda.is_available():
         amp = torch.cuda.amp.autocast(enabled=False)
@@ -135,14 +132,6 @@ def run_inference(dataset, model, epochs, profiling, bf16, use_compile,
             with torch_profile():
                 inference(model, data)
 
-        if use_compile:
-            print("Using torch.compile")
-            compiled_model = torch_geometric.compile(model)
-            inference(compiled_model, data)
-            inference(compiled_model, data)
-            with timeit():
-                inference(compiled_model, data)
-
 
 def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
         inference, profiling, bf16, use_compile, permute_masks=None,

benchmark/kernel/main_performance.py

@@ -79,7 +79,8 @@ def run_train():
 
             model = Model(dataset, num_layers, hidden).to(device)
             optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
-
+            if args.compile:
+                model = torch_geometric.compile(model)
             loss_list = []
             acc_list = []
             for epoch in range(1, args.epochs + 1):
@@ -106,14 +107,6 @@ def run_train():
                 rename_profile_file(model_name, dataset_name, str(num_layers),
                                     str(hidden), 'train')
 
-            if args.compile:
-                print("Using torch.compile")
-                compiled_model = torch_geometric.compile(model)
-                eval_acc(compiled_model, val_loader)
-                eval_acc(compiled_model, val_loader)
-                with timeit():
-                    eval_acc(compiled_model, val_loader)
-
 
 @torch.no_grad()
 def run_inference():
@@ -126,7 +119,8 @@ def run_inference():
             print(f'{dataset_name} - {model_name}- {num_layers} - {hidden}')
 
             model = Model(dataset, num_layers, hidden).to(device)
-
+            if args.compile:
+                model = torch_geometric.compile(model)
             with amp:
                 for epoch in range(1, args.epochs + 1):
                     if epoch == args.epochs:
@@ -142,14 +136,6 @@ def run_inference():
                                         str(num_layers), str(hidden),
                                         'inference')
 
-                if args.compile:
-                    print("Using torch.compile")
-                    compiled_model = torch_geometric.compile(model)
-                    inference_run(compiled_model, test_loader, args.bf16)
-                    inference_run(compiled_model, test_loader, args.bf16)
-                    with timeit():
-                        inference_run(compiled_model, test_loader, args.bf16)
-
 
 if not args.inference:
     run_train()

benchmark/points/train_eval.py

@@ -15,6 +15,8 @@ def run_train(train_dataset, test_dataset, model, epochs, batch_size,
               use_compile, lr, lr_decay_factor, lr_decay_step_size,
               weight_decay):
     model = model.to(device)
+    if use_compile:
+        model = torch_geometric.compile(model)
     optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
 
     train_loader = DataLoader(train_dataset, batch_size, shuffle=True)
@@ -41,19 +43,13 @@ def run_train(train_dataset, test_dataset, model, epochs, batch_size,
             for param_group in optimizer.param_groups:
                 param_group['lr'] = lr_decay_factor * param_group['lr']
 
-    if use_compile:
-        print("Using torch.compile")
-        compiled_model = torch_geometric.compile(model)
-        test(compiled_model, test_loader, device)
-        test(compiled_model, test_loader, device)
-        with timeit():
-            test(compiled_model, test_loader, device)
-
 
 @torch.no_grad()
 def run_inference(test_dataset, model, epochs, batch_size, profiling, bf16,
                   use_compile):
     model = model.to(device)
+    if use_compile:
+        model = torch_geometric.compile(model)
     test_loader = DataLoader(test_dataset, batch_size, shuffle=False)
 
     if torch.cuda.is_available():
@@ -74,14 +70,6 @@ def run_inference(test_dataset, model, epochs, batch_size, profiling, bf16,
             with torch_profile():
                 inference(model, test_loader, device, bf16)
 
-        if use_compile:
-            print("Using torch.compile")
-            compiled_model = torch_geometric.compile(model)
-            inference(compiled_model, test_loader, device, bf16)
-            inference(compiled_model, test_loader, device, bf16)
-            with timeit():
-                inference(compiled_model, test_loader, device, bf16)
-
 
 def run(train_dataset, test_dataset, model, epochs, batch_size, lr,
         lr_decay_factor, lr_decay_step_size, weight_decay, inference,