update gnn

course_UvA-DL/06-graph-neural-networks/GNN_overview.py

@@ -634,26 +634,27 @@ def test_step(self, batch, batch_idx):
 # Additionally to the Lightning module, we define a training function below.
 # As we have a single graph, we use a batch size of 1 for the data loader and share the same data loader for the train,
 # validation, and test set (the mask is picked inside the Lightning module).
-# Besides, we set the argument `progress_bar_refresh_rate` to zero as it usually shows the progress per epoch,
+# Besides, we set the argument `enable_progress_bar` to False as it usually shows the progress per epoch,
 # but an epoch only consists of a single step.
 # If you have downloaded the pre-trained models in the beginning of the tutorial, we load those instead of training from scratch.
 # Finally, we test the model and return the results.
 
 
 # %%
 def train_node_classifier(model_name, dataset, **model_kwargs):
-    pl.seed_everything(42)
+    L.seed_everything(42)
     node_data_loader = geom_data.DataLoader(dataset, batch_size=1)
 
     # Create a PyTorch Lightning trainer
     root_dir = os.path.join(CHECKPOINT_PATH, "NodeLevel" + model_name)
     os.makedirs(root_dir, exist_ok=True)
-    trainer = pl.Trainer(
+    trainer = L.Trainer(
         default_root_dir=root_dir,
         callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc")],
-        gpus=AVAIL_GPUS,
+        accelerator="auto",
+        devices=AVAIL_GPUS,
         max_epochs=200,
-        progress_bar_refresh_rate=0,
+        enable_progress_bar=False,
     )  # 0 because epoch size is 1
     trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need
 
@@ -663,7 +664,7 @@ def train_node_classifier(model_name, dataset, **model_kwargs):
         print("Found pretrained model, loading...")
         model = NodeLevelGNN.load_from_checkpoint(pretrained_filename)
     else:
-        pl.seed_everything()
+        L.seed_everything()
         model = NodeLevelGNN(
             model_name=model_name, c_in=dataset.num_node_features, c_out=dataset.num_classes, **model_kwargs
         )