Feature/virtual nodes

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 core.python.*
+core.python.*
 
 # C extensions
 *.so
@@ -179,8 +180,14 @@ test*.ipynb
 
 checkpoints/
 
+# local configs
+configs/examples/cgcnn_vn/
+configs/examples/cgcnn_vn_hg/
+configs/generated
+
 # misc
 .flake8
 .pylintrc
 **/wandb/
 *.out
+results/

.vscode/settings.json

@@ -0,0 +1,9 @@
+{
+    "python.analysis.typeCheckingMode": "off",
+    "python.linting.flake8Enabled": true,
+    "python.linting.enabled": true,
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter"
+    },
+    "python.formatting.provider": "none"
+}

configs/config.yml

@@ -4,23 +4,27 @@ task:
   run_mode: "train"
   identifier: "my_train_job"
   parallel: False
-  # seed=0 means random initalization
+  use_wandb: True
+  wandb_entity: "fung-lab"
+  wandb_project: "cgcnn_vn"
+
   seed: 12345678
   # Defaults to run directory if not specified
   save_dir: 
   continue_job: False
   load_training_state: False
   # Path to the checkpoint file
-  checkpoint_path: 
+  checkpoint_dir: 
   write_output: True  
 
 model:
   name: CGCNN  
   save_model: True
-  model_path: "my_model.pth"
+  model_path: "cgcnn.pth"
   edge_steps: 50
   self_loop: True
   # model attributes
+  # model attributes
   dim1: 100
   dim2: 150
   pre_fc_count: 1
@@ -35,14 +39,17 @@ model:
 
 optim:
   max_epochs: 40
+  max_epochs: 250
+  max_checkpoint_epochs: 0
   max_checkpoint_epochs: 0
   lr: 0.002
   # Either custom or from torch.nn.functional library. If from torch, loss_type is TorchLossWrapper
+  # Either custom or from torch.nn.functional library. If from torch, loss_type is TorchLossWrapper
   loss:
     loss_type: "TorchLossWrapper"
     loss_args: {"loss_fn": "l1_loss"}
 
-  batch_size: 100
+  batch_size: 64
   optimizer:
     optimizer_type: "AdamW"
     optimizer_args: {}
@@ -82,8 +89,9 @@ dataset:
   node_representation: "onehot"
   additional_attributes: []
   # Print out processing info
+  # Print out processing info
   verbose: True
-  # Index of target column in targets.csv
+  use_virtual_nodes: False
   # graph specific settings
   preprocess_params:
     cutoff_radius : 8.0

configs/config_templates/cgcnn_2dnpj.yml

@@ -0,0 +1,140 @@
+trainer: property
+
+task:
+  # run_mode: train
+  identifier: "cgcnn_vn_2dnpj"
+  reprocess: False
+  run_id: ""
+  parallel: True
+  device: "cuda:0"
+  seed: 0
+  # seed=0 means random initalization
+  write_output: True
+  parallel: True
+  # Training print out frequency (print per n number of epochs)
+  verbosity: 1
+  wandb:
+    use_wandb: True
+    wandb_entity: "fung-lab"
+    wandb_project: "cgcnn_vn_new"
+    notes: "try 2d dataset"
+    tags: ["test"]
+    track_params:
+      - "model.hyperparams.mp_pattern"
+      - "model.hyperparams.pool"
+      - "model.hyperparams.virtual_pool"
+      - "model.hyperparams.gc_count"
+      - "optim.lr"
+      - "optim.batch_size"
+      - "optim.max_epochs"
+      - "dataset.preprocess_params.num_offsets"
+      - "dataset.preprocess_params.edge_calc_method"
+      - "dataset.preprocess_params.all_neighbors"
+      - "dataset.preprocess_params.edge_steps"
+      - "dataset.preprocess_params.use_degree"
+    log_artifacts:
+      - "/global/cfs/projectdirs/m3641/Sidharth/MatDeepLearn_dev/matdeeplearn/models/cgcnn_vn.py"
+    metadata:
+      architecture: "CGCNN_VN"
+      cluster: "fung-cluster"
+      dataset: "2DNPJ_data"
+    sweep:
+      parallel: True
+      do_sweep: False # ignore rest of config if False
+      system: "phoenix_slurm" # one of "local", "phoenix_slurm"
+      job_config: "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/configs/jobs/phoenix_slurm.yml"
+      count: 3
+      sweep_file: "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/configs/sweeps/cgcnn_vn_sweep_d1.yml"
+
+model:
+  name: CGCNN_VN
+  load_model: False
+  save_model: True
+  model_path: "cgcnn_vn.pth"
+  # model hyperparams
+  hyperparams:
+    edge_steps: 25
+    self_loop: True
+    dim1: 100
+    dim2: 150
+    atomic_intermediate_layer_resolution: 0
+    pre_fc_count: 1
+    gc_count: 4
+    post_fc_count: 3
+    pool: "global_mean_pool" # pooling reduction scheme
+    virtual_pool: "AtomicNumberPooling" # pooling method
+    pool_choice: "virtual" # whether to use virtual or real nodes or both in RealVirtualPooling
+    mp_pattern: ["rr", "rv"]
+    pool_order: "early"
+    batch_norm: True
+    batch_track_stats: True
+    act_fn: "relu"
+    act_nn: "ReLU"
+    dropout_rate: 0.0
+
+optim:
+  max_epochs: 250
+  lr: 0.002
+  loss:
+    loss_type: "TorchLossWrapper"
+    loss_args: {"loss_fn": "l1_loss"}
+  batch_size: 64
+  optimizer:
+    optimizer_type: "AdamW"
+    optimizer_args: {}
+  scheduler:
+    scheduler_type: "ReduceLROnPlateau"
+    scheduler_args: {"mode":"min", "factor":0.8, "patience":10, "min_lr":0.00001, "threshold":0.0002}
+
+dataset:
+  processed: False # if False, need to preprocessor data and generate .pt file
+  force_preprocess: False
+  num_examples: 0 # set to 0 when using full dataset, else will take the first "num_examples" examples
+  # Path to data files
+  # src: "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/2D_data_npj/raw"
+  src: "/global/cfs/projectdirs/m3641/Sidharth/datasets/2D_data_npj/raw"
+  # target_path: "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/2D_data_npj/targets.csv"
+  target_path: "/global/cfs/projectdirs/m3641/Sidharth/datasets/2D_data_npj/targets.csv"
+  # pt_path: "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/2D_data_npj/processed/ocp"
+  pt_path: "/global/cfs/projectdirs/m3641/Sidharth/datasets/2D_data_npj/processed/ocp"
+  # transforms
+  transforms:
+    - name: GetY
+      args:
+        index: 0
+      otf: False
+    - name: VirtualNodeGeneration
+      args:
+        virtual_box_increment: 3
+      otf: False
+    - name: VirtualEdgeGeneration
+      args:
+        attrs: ["rr", "rv"]
+        rr_cutoff: 12.0
+        rv_cutoff: 12.0
+      otf: False
+      batch: True
+  # use for passing into global config
+  # one of MDL, ASE, OCP
+  use_sweep_params: False
+  apply_pre_transform_processing: False
+  # use again for passing into global config
+  data_format: "json"
+  node_representation: "onehot"
+  additional_attributes: []
+  # Print out processing info
+  verbose: True
+  # graph specific settings: preprocessing hyperparams
+  preprocess_params:
+    cutoff_radius : 5.0
+    n_neighbors : 250
+    process_batch_size : 100
+    edge_calc_method: "ocp"
+    num_offsets: 1
+    edge_steps : 25
+    all_neighbors: True
+    use_degree: False
+  # Ratios for train/val/test split out of a total of 1
+  train_ratio: 0.8
+  val_ratio: 0.05
+  test_ratio: 0.15

configs/config_templates/cgcnn_han_hmof_5k.yml

@@ -0,0 +1,152 @@
+trainer: property
+
+task:
+  # run_mode: train
+  identifier: "cgcnn_han_vn_post_hmof"
+  reprocess: False
+  run_id: ""
+  parallel: True
+  device: "cuda:0"
+  seed: 0
+  # seed=0 means random initalization
+  write_output: True
+  parallel: True
+  # Training print out frequency (print per n number of epochs)
+  verbosity: 1
+  wandb:
+    use_wandb: True
+    wandb_entity: "fung-lab"
+    wandb_project: "cgcnn_vn_new"
+    notes: ""
+    tags: ["test"]
+    track_params:
+      - "model.hyperparams.mp_pattern"
+      - "model.hyperparams.pool"
+      - "model.hyperparams.virtual_pool"
+      - "model.hyperparams.gc_count"
+      - "model.hyperparams.attn_heads"
+      - "optim.lr"
+      - "optim.batch_size"
+      - "optim.max_epochs"
+      - "dataset.preprocess_params.num_offsets"
+      - "dataset.preprocess_params.edge_calc_method"
+      - "dataset.preprocess_params.all_neighbors"
+      - "dataset.preprocess_params.edge_steps"
+      - "dataset.preprocess_params.use_degree"
+    log_artifacts:
+      # - "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/matdeeplearn/models/cgcnn_vn.py"
+      # - "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/MatDeepLearn_dev/matdeeplearn/models/cgcnn_vn.py"
+      - "/global/cfs/projectdirs/m3641/Sidharth/MatDeepLearn_dev/matdeeplearn/models/cgcnn_vn.py"
+    metadata:
+      architecture: "CGCNN_VN"
+      cluster: "fung-cluster"
+      dataset: "hMOF-5K"
+    sweep:
+      parallel: True
+      do_sweep: False # ignore rest of config if False
+      system: "phoenix_slurm" # one of "local", "phoenix_slurm"
+      job_config: "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/configs/jobs/phoenix_slurm.yml"
+      count: 3
+      sweep_file: "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/configs/sweeps/cgcnn_vn_sweep_d1.yml"
+      # sweep_file: "/nethome/sbaskaran31/projects/Sidharth/MatDeepLearn_dev/configs/cgcnn_vn_sweep.yml"
+
+model:
+  name: CGCNN_HAN_VN
+  load_model: False
+  save_model: True
+  model_path: "cgcnn_han_vn.pth"
+  # model hyperparams
+  hyperparams:
+    edge_steps: 25
+    self_loop: True
+    dim1: 100
+    dim2: 150
+    atomic_intermediate_layer_resolution: 0
+    pre_fc_count: 1
+    gc_count: 4
+    post_fc_count: 3
+    attn_heads: 6
+    pool: "global_mean_pool" # pooling reduction scheme
+    virtual_pool:
+      virtual_pool_name: "RealVirtualAttention" # pooling method
+      args:
+        embed_dim: 100
+        attn_size: 128
+    mp_pattern: ["rr", "rv"]
+    pool_order: "early"
+    batch_norm: True
+    batch_track_stats: True
+    act_fn: "relu"
+    act_nn: "ReLU"
+    dropout_rate: 0.0
+
+optim:
+  max_epochs: 250
+  lr: 0.002
+  loss:
+    loss_type: "TorchLossWrapper"
+    loss_args: {"loss_fn": "l1_loss"}
+  batch_size: 64
+  optimizer:
+    optimizer_type: "AdamW"
+    optimizer_args: {}
+  scheduler:
+    scheduler_type: "ReduceLROnPlateau"
+    scheduler_args: {"mode":"min", "factor":0.8, "patience":10, "min_lr":0.00001, "threshold":0.0002}
+
+dataset:
+  processed: False # if False, need to preprocessor data and generate .pt file
+  force_preprocess: False
+  num_examples: 0 # set to 0 when using full dataset, else will take the first "num_examples" examples
+  # Path to data files
+  # src: "/nethome/sbaskaran31/projects/Sidharth/hMOF/raw_5k/data.json"
+  # src: "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/hMOF/raw_5k/raw_5k/data.json"
+  src: "/global/cfs/projectdirs/m3641/Shared/Materials_datasets/hMOF/raw_5k/data.json"
+  target_path: ""
+  # pt_path: "/nethome/sbaskaran31/projects/Sidharth/hMOF/raw_5k/ocp"
+  # pt_path: "/storage/home/hcoda1/9/sbaskaran31/p-vfung3-0/hMOF/raw_5k/raw_5k/ocp"
+  pt_path: "/global/cfs/projectdirs/m3641/Shared/Materials_datasets/hMOF/raw_5k/ocp"
+  # transforms
+  transforms:
+    - name: GetY
+      args:
+        index: 5 # methane adsorption uptake
+      otf: False
+    - name: VirtualNodeGeneration
+      args:
+        virtual_box_increment: 3
+        method: "ase"
+      otf: False
+    - name: VirtualEdgeGeneration
+      args:
+        attrs: ["rr", "rv"]
+        rr_cutoff: 12
+        rv_cutoff: 12
+        vr_cutoff: 5.0
+        vv_cutoff: 5.0
+      otf: False
+      batch: True
+  # use for passing into global config
+  # one of MDL, ASE, OCP
+  use_sweep_params: False
+  apply_pre_transform_processing: False
+  # use again for passing into global config
+  data_format: "json"
+  node_representation: "onehot"
+  additional_attributes: []
+  # Print out processing info
+  verbose: True
+  # graph specific settings: preprocessing hyperparams
+  preprocess_params:
+    cutoff_radius : 5.0
+    n_neighbors : 250
+    process_batch_size : 50
+    edge_calc_method: "ocp"
+    num_offsets: 1
+    edge_steps : 25
+    all_neighbors: True
+    use_degree: False
+  # Ratios for train/val/test split out of a total of 1
+  train_ratio: 0.8
+  val_ratio: 0.05
+  test_ratio: 0.15