minor tweaks to testing

.gitignore

@@ -91,6 +91,8 @@ ipython_config.py
 #   install all needed dependencies.
 #Pipfile.lock
 
+.vscode*
+
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 __pypackages__/
 

configs/datamodule/default.yaml

@@ -7,7 +7,7 @@ data_conf:
   jet_type: t
   data_dir: ${paths.data_dir}
   num_particles: 30
-  split_fraction: [0.7, 0.0, 0.3]
+  split_fraction: [0.9, 0.0, 0.1]
   particle_features:
     - etarel
     - phirel
@@ -20,5 +20,5 @@ data_conf:
 loader_kwargs:
   pin_memory: true
   batch_size: 256
-  num_workers: 2
+  num_workers: 16
   drop_last: True

configs/model/default.yaml

@@ -22,26 +22,26 @@ normaliser_config:
 # Full transformer encoder model
 trans_enc_config:
   te_config:
-    model_dim: 64
-    num_layers: 2
+    model_dim: 128
+    num_layers: 4
     mha_config:
       num_heads: 8
     dense_config:
-      hddn_dim: 64
+      hddn_dim: 256
       act_h: lrlu
       nrm: layer
   node_embd_config:
-    num_blocks: 1
-    hddn_dim: 64
+    hddn_dim: 256
     act_h: lrlu
     nrm: layer
   outp_embd_config:
-    hddn_dim: 64
+    hddn_dim: 256
     act_h: lrlu
     nrm: layer
   ctxt_embd_config:
     outp_dim: 64
-    hddn_dim: 64
+    hddn_dim: 128
+    num_blocks: 2
     act_h: lrlu
     nrm: layer
 

configs/paths/default.yaml

@@ -1,6 +1,6 @@
 # output dir will be created by hydra
-data_dir: /srv/beegfs/scratch/groups/rodem/datasets/jetnet/
-output_dir: /srv/beegfs/scratch/groups/rodem/jet_diffusion/checkpoints
+data_dir: /home/matthew/Documents/data/jetnet/
+output_dir: /home/matthew/Documents/saved_networks
 
 # Interpolated
 root_dir: ${oc.env:PROJECT_ROOT}

configs/trainer/default.yaml

@@ -7,7 +7,7 @@ accelerator: auto
 devices: 1
 gradient_clip_val: 5
 precision: 32
-check_val_every_n_epoch: 1
+check_val_every_n_epoch: 10
 
 # Interpolated
 default_root_dir: ${paths.full_path}

scripts/train.py

@@ -23,6 +23,7 @@
     version_base=None, config_path=str(root / "configs"), config_name="train.yaml"
 )
 def main(cfg: DictConfig) -> None:
+
     log.info("Setting up full job config")
     if cfg.full_resume:
         cfg = reload_original_config(cfg)

src/models/pc_jedi.py

@@ -178,7 +178,7 @@ def validation_step(self, sample: tuple, batch_idx: int) -> None:
         )
 
         # Add to the collection of the validaiton outputs
-        self.val_outs.append(to_np(outputs), to_np(sample))
+        self.val_outs.append((to_np(outputs), to_np(sample)))
 
     def on_validation_epoch_end(self) -> None:
         """At the end of the validation epoch, calculate and log the metrics