Add Diffusion Graph Network

3rd_party/gnn/dist-dgn/client.py

@@ -0,0 +1,16 @@
+import os
+import sys
+from typing import Optional, Union, Tuple
+import logging
+from omegaconf import DictConfig
+import numpy as np
+from time import sleep, perf_counter
+
+log = logging.getLogger(__name__)
+
+
+class OnlineClient:
+    """Class for the online training client"""
+
+    def __init__(self) -> None:
+        self.client = None

3rd_party/gnn/dist-dgn/conf/config.yaml

@@ -0,0 +1,95 @@
+# @package _global_
+verbose: False
+timers: False
+postprocess: False
+statistics: False
+seed: 12
+backend: null
+num_threads: 1
+logfreq: 10
+ckptfreq: 500
+batch_size: 1
+val_batch_size: 1
+precision: fp32
+fp16_allreduce: False
+restart: False
+master_addr: none
+master_port: 2345
+device_skip: 0 # temporary workaround to skip GPU on a node
+
+# DGN specific parameters
+prediction_type: "epsilon" # {epsilon, x0, v}
+num_diffusion_steps: 100
+emb_width: 128
+diffusion_process_schedule: "linear" # {linear, cosine}
+diffusion_step_sampler: "uniform" # {uniform, adaptive_exponential}
+num_gen_samples: 1
+input_node_features: 3
+learnable_variance: False
+loss_weighting: "uniform" # {uniform, min_snr}
+min_snr_gamma: 5.0 # gamma for min-SNR weighting (only used when loss_weighting=min_snr)
+activation_checkpointing: False 
+cond_node_features: False
+
+# Modeling task
+model_task: "train"
+consistency: True
+
+# learning rate schedule and training steps
+phase1_steps: 100 
+phase2_steps: 0
+phase3_steps: 0
+lr_phase12: 0.0001
+lr_phase23: 0.0001
+
+# model arch properties
+model_name: "gnn" # {gnn, graph_transformer}
+mlp_hidden_channels: 32
+n_mlp_hidden_layers: 5
+n_messagePassing_layers: 4
+n_transformer_layers: 4
+num_heads: 4
+layer_norm: False
+dropout_rate: 0.0
+
+# Halo swap mode 
+halo_swap_mode : none
+
+# transform directions (enforces periodicity in graph)
+transform_x: false
+transform_y: false
+transform_z: false
+
+# Correctness validation
+target_loss: 0
+
+# plotting the connecivity (can take some time)
+plot_connectivity : False
+
+# specify path to gnn_outputs here (case directory)  
+gnn_outputs_path: ${work_dir}/gnn_outputs_poly_7/
+traj_data_path: ${work_dir}/traj_poly_1/
+
+# Online training
+online: False
+
+hydra:
+  job:
+    chdir: false
+
+    #defaults:
+    #  - override hydra/hydra_logging: colorlog
+    #  - override hydra/job_logging: colorlog
+
+# path to original working directory
+# hydra hijacks working directory by changing it to the new log directory
+# so its useful to have this path as a special variable
+# https://hydra.cc/docs/next/tutorials/basic/running_your_app/working_directory
+work_dir: ${hydra:runtime.cwd}
+# path to folder with data
+data_dir: ${work_dir}/datasets/
+# path to folder for checkpointing
+ckpt_dir: ${work_dir}/ckpt/ 
+# path to saved model directory
+model_dir: ${work_dir}/saved_models/ 
+