update ptreain config

Author: Oxer11

config/pretrain/angle_gearnet_edge.yaml

@@ -1,11 +1,11 @@
-output_dir: ~/scratch/protein_outputs
+output_dir: ~/scratch/protein_output
 
 dataset:
   class: AlphaFoldDB
   path: ~/scratch/protein-datasets/alphafold
   species_start: 0
   species_end: 22
-  # species_id: 0
+  # species_id: 3
   # split_id: 1
   transform:
     class: Compose

config/pretrain/dihedral_gearnet_edge.yaml

@@ -1,4 +1,4 @@
-output_dir: ~/scratch/protein_outputs
+output_dir: ~/scratch/protein_output
 
 dataset:
   class: AlphaFoldDB

config/pretrain/dis_gearnet_edge.yaml

@@ -1,4 +1,4 @@
-output_dir: ~/scratch/protein_outputs
+output_dir: ~/scratch/protein_output
 
 dataset:
   class: AlphaFoldDB

config/pretrain/mc_ESM_gearnet.yaml renamed to config/pretrain/mc_gearnet_edge.yaml

@@ -1,11 +1,11 @@
-output_dir: ~/scratch/protein_outputs
+output_dir: ~/scratch/protein_output
 
 dataset:
   class: AlphaFoldDB
   path: ~/scratch/protein-datasets/alphafold
   species_start: 0
   species_end: 22
-  # species_id: 0
+  # species_id: 3
   # split_id: 1
   transform:
     class: ProteinView

util.py

@@ -120,6 +120,8 @@ def build_downstream_solver(cfg, dataset):
         solver.load(cfg.checkpoint)
 
     if cfg.get("model_checkpoint") is not None:
+        if comm.get_rank() == 0:
+            logger.warning("Load checkpoint from %s" % cfg.model_checkpoint)
         cfg.model_checkpoint = os.path.expanduser(cfg.model_checkpoint)
         model_dict = torch.load(cfg.model_checkpoint, map_location=torch.device('cpu'))
         task.model.load_state_dict(model_dict)
@@ -137,4 +139,4 @@ def build_pretrain_solver(cfg, dataset):
     optimizer = core.Configurable.load_config_dict(cfg.optimizer)
     solver = core.Engine(task, dataset, None, None, optimizer, **cfg.engine)
 
-    return solver
+    return solver