Merge pull request #35 from joschkabirk/dev

JetClass improvements

configs/callbacks/jetclass.yaml

@@ -8,7 +8,7 @@ defaults:
   # - device_stats_monitor.yaml
   - rich_progress_bar.yaml
   - ema.yaml
-  - jetclass_eval_test.yaml
+  # - jetclass_eval_test.yaml  # ignore for now since crashes with sigterm 15
   - jetclass_eval.yaml
   - _self_
 

configs/data/jetclass.yaml

@@ -20,8 +20,11 @@ used_jet_types: null # null means all jet types
 additional_part_features: null # null means no additional part features
 
 # files and jet types to use
-data_dir: /beegfs/desy/user/birkjosc/datasets/jetclass_h5_more_part_features_including_binary
+data_dir: /beegfs/desy/user/birkjosc/datasets/jetclass_h5_all_features_5M_training_jets/
 filename_dict:
   train: ${data.data_dir}/train_100M/merged_standardized.h5
   val: ${data.data_dir}/val_5M/merged_standardized.h5
   test: ${data.data_dir}/test_20M/merged_standardized.h5
+# the file below contains the masks and m/pt/nparticles generated by another model
+# --> use for generating new point clouds
+conditioning_gen_filename: /beegfs/desy/user/birkjosc/datasets/jetclass_cond_gen_data/flow_2hidden_layers_512hidden_features_200_000steps_0.99998gamma_2_000_000jets.h5

configs/data/jetclass_dev.yaml

@@ -20,8 +20,11 @@ used_jet_types: null # null means all jet types
 additional_part_features: null # null means no additional part features
 
 # files and jet types to use
-data_dir: /beegfs/desy/user/birkjosc/datasets/jetclass_h5_more_part_features_including_binary_dev_shuffled
+data_dir: /beegfs/desy/user/birkjosc/datasets/jetclass_h5_all_features_5M_training_jets_dev/
 filename_dict:
   train: ${data.data_dir}/train_100M/merged_standardized.h5
   val: ${data.data_dir}/val_5M/merged_standardized.h5
   test: ${data.data_dir}/test_20M/merged_standardized.h5
+# the file below contains the masks and m/pt/nparticles generated by another model
+# --> use for generating new point clouds
+conditioning_gen_filename: /beegfs/desy/user/birkjosc/datasets/jetclass_cond_gen_data/flow_2hidden_layers_512hidden_features_200_000steps_0.99998gamma_2_000_000jets.h5

configs/experiment/jetclass_cond.yaml

@@ -19,20 +19,23 @@ tags: ["flow_matching", "JetClass", "cond"]
 
 run_note: ""
 
-seed: 12345
+seed: 160397
 
 trainer:
   min_epochs: 1
   max_epochs: 1000
   gradient_clip_val: 0.5
 
+# setting load_weights_from will load the weights from the given checkpoint, but start training from scratch
+# load_weights_from: /beegfs/desy/user/birkjosc/epic-fm/logs/jetclass/runs/2023-08-24_15-36-42/checkpoints/epoch_160_loss_40.79013-EMA.ckpt
+
 model:
   num_particles: 128
   global_cond_dim: 10 # needs to be calculated when using conditioning (= number of jet types + additional conditioning variables)
   local_cond_dim: 0
   features: 13  # = 3 + number of `additional_part_features` (see below)
-  layers: 10 # default is 6
-  hidden_dim: 200 # default is 128
+  layers: 13 # default is 6
+  hidden_dim: 250 # default is 128
   latent: 15 # default is 10
   scheduler:
     warmup: 50
@@ -53,6 +56,7 @@ data:
   conditioning_mass: False
   conditioning_num_particles: False
   conditioning_jet_type: True
+  conditioning_jet_type_all: True  # if true, all jet types are used for conditioning even if trained on only some of them (see `used_jet_types` above)
   additional_part_features:
     - part_d0val
     - part_d0err
@@ -73,11 +77,11 @@ callbacks:
     save_ema_weights_in_callback_state: True
     evaluate_ema_weights_instead: True
   jetclass_eval:
-    every_n_epochs: 50 # evaluate every n epochs
-    additional_eval_epochs: [1, 30, 75] # evaluate at these epochs as well
+    every_n_epochs: 100 # evaluate every n epochs
+    additional_eval_epochs: [30, 75] # evaluate at these epochs as well
     num_jet_samples: 50_000 # jet samples to generate
-  jetclass_eval_test:
-    num_jet_samples: 200_000 # jet samples to generate
+  # jetclass_eval_test:
+  #   num_jet_samples: 50_000 # jet samples to generate
 
   #early_stopping:
   #  monitor: "val/loss"

configs/experiment/jetclass_dev.yaml

@@ -15,6 +15,8 @@ defaults:
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters
 
+# add here checkpoint to continue training
+# ckpt_path: /beegfs/desy/user/birkjosc/epic-fm/logs/jetclass_cond_jettype/runs/2023-08-10_16-36-26/checkpoints/last-EMA.ckpt
 
 
 tags: ["flow_matching", "JetClass", "dev", "debug"]
@@ -28,14 +30,17 @@ trainer:
   max_epochs: 4
   gradient_clip_val: 0.5
 
+load_weights_from: /beegfs/desy/user/birkjosc/epic-fm/logs/jetclass/runs/2023-08-23_23-53-11/checkpoints/last-EMA.ckpt
+
 model:
   num_particles: 128
   global_cond_dim: 10 # needs to be calculated when using conditioning (= number of jet types + additional conditioning variables)
   local_cond_dim: 0
-  features: 4
+  features: 3
 
 data:
   # preprocessing
+  # used_jet_types: [Tbqq] # null means all jet types
   loss_per_jettype: true
   number_of_used_jets: 30_000
   number_of_used_jets_val: 10_000
@@ -47,9 +52,10 @@ data:
   conditioning_mass: False
   conditioning_num_particles: False
   conditioning_jet_type: True
-  additional_part_features:
+  conditioning_jet_type_all: true
+  # additional_part_features:
     # - part_energy
-    - part_d0val
+    # - part_d0val
     # - part_d0err
     # - part_dzval
     # - part_dzerr
@@ -81,5 +87,5 @@ logger:
     group: "jetclass_flow_matching_dev"
     name: ${task_name}
   comet:
-    experiment_name: null
+    experiment_name: ${task_name}
     project_name: "flow-matching"