Merge pull request #38 from ewencedr/cedric-dev

Transformers

configs/data/classifier_data.yaml

@@ -7,4 +7,7 @@ pin_memory: False
 gen_jet: "both"
 ref_jet: "both"
 idealized: False
-use_shuffled_data: False
+use_shuffled_data: False
+use_nonrel_data: False
+n_signal: 0
+n_background: 100_000

configs/experiment/lhco/epic_classifier.yaml

@@ -35,8 +35,8 @@ callbacks:
     verbose: true
 
 data:
-  #gendatafile: "lhco-xy-10layer-256latent-logpt-new2m_sr-midpoint-300_shuffled"
-  gendatafile: "lhco-xy-10layer-256latent-logpt-nnew_sr-midpoint-300"
+  gendatafile: "lhco-xy-10layer-256latent-logpt-new2m_sr-midpoint-300_shuffled"
+  #gendatafile: "lhco-xy-10layer-256latent-logpt-nnew_sr-midpoint-300"
   #gendatafile: "FPCD_LHCO_SR_shuffled"
 
 task_name: "epic_classifier"

configs/experiment/lhco/lhco.yaml → configs/experiment/lhco/lhco_bigPC.yaml

@@ -6,39 +6,50 @@
 defaults:
   - override /data: lhco.yaml
   - override /model: flow_matching.yaml
-  - override /callbacks: lhco.yaml
-  - override /trainer: gpu.yaml
+  - override /callbacks: /lhco/lhco_bruteforce.yaml
+  - override /trainer: ddp.yaml
 
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters
 
 
 
-tags: ["flow_matching", "LHCO", "uncond"]
+tags: ["flow_matching", "LHCO", "bigPC"]
 
 seed: 12345
 
 trainer:
-  min_epochs: 100
-  max_epochs: 1000
+  min_epochs: 150
+  max_epochs: 3000
   gradient_clip_val: 0.5
 
+
 model:
-  num_particles: 560
-  global_cond_dim: 0 # needs to be calculated when using conditioning
-  local_cond_dim: 0
+  num_particles: 558
+  global_cond_dim: 10 # needs to be calculated when using conditioning
+  local_cond_dim: 10
+  layers: 8
+  latent: 256
+  hidden_dim: 256
   scheduler:
     warmup: ${trainer.min_epochs}
     max_iters: ${trainer.max_epochs}
 
 data:
+  batch_size: 128
   normalize: True
-  conditioning: False
+  conditioning: True
   relative_coords: False
   jet_type: "all_one_pc"
   num_particles: ${model.num_particles}
+  use_all_data: False
+  shuffle_data: False
   val_fraction: 0.05
   test_fraction: 0.35
+  file_suffix_processed_data: "_masssorted"
+  log_pt: False
+  pt_standardization: False
+  multiplicity_conditioning: True
 
 callbacks:
   ema:
@@ -48,8 +59,8 @@ callbacks:
     save_ema_weights_in_callback_state: True
     evaluate_ema_weights_instead: True
 
-  jetnet_eval:
-    every_n_epochs: 100 # evaluate every n epochs
+  lhco_bruteforce_eval:
+    every_n_epochs: 300 # evaluate every n epochs
     num_jet_samples: -1 # jet samples to generate
     model_name: "lhco_flow_matching"
     log_epoch_zero: True

configs/experiment/lhco/lhco_both_jets2.yaml

@@ -28,7 +28,7 @@ model:
   global_cond_dim: 5 # needs to be calculated when using conditioning
   local_cond_dim: 5
   layers: 12
-  hidden_dim: 150
+  hidden_dim: 256
   latent: 256
   #scheduler:
   #  warmup: ${trainer.min_epochs}

configs/experiment/lhco/lhco_both_jets_crossattention.yaml

@@ -0,0 +1,119 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=fm_tops
+
+defaults:
+  - override /data: lhco.yaml
+  - override /model: fm_droid_crossattention.yaml
+  - override /callbacks: /lhco/lhco.yaml
+  - override /trainer: ddp.yaml
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+
+
+tags: ["flow_matching", "LHCO", "cond", "both_jet", "crossattention"]
+
+seed: 12345
+compile: False
+
+trainer:
+  min_epochs: 250
+  max_epochs: 5000
+  gradient_clip_val: 0.5
+
+model:
+  num_particles: 279
+  global_cond_dim: 5 # needs to be calculated when using conditioning
+  scheduler:
+    warmup: ${trainer.min_epochs}
+    max_iters: ${trainer.max_epochs}
+  #optimizer:
+  #  _target_: torch.optim.AdamW
+  #  _partial_: true
+  #  lr: 0.0003
+  #  weight_decay: 0.00005
+
+  #scheduler:
+  #  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  #  _partial_: true
+  #  T_max: ${trainer.max_epochs}
+
+data:
+  batch_size: 128
+  normalize: True
+  normalize_sigma: 1
+  conditioning: True
+  relative_coords: True
+  jet_type: "all"
+  num_particles: ${model.num_particles}
+  val_fraction: 0.05
+  test_fraction: 0.35
+  use_all_data: False
+  shuffle_data: False
+  file_suffix_processed_data: "_masssorted"
+  log_pt: True
+  pt_standardization: False
+  multiplicity_conditioning: True
+
+callbacks:
+  ema:
+    decay: 0.999
+    apply_ema_every_n_steps: 1
+    start_step: 0
+    save_ema_weights_in_callback_state: True
+    evaluate_ema_weights_instead: True
+
+  early_stopping: null
+
+  #lhco_eval:
+  #  every_n_epochs: 300 # evaluate every n epochs
+  #  num_jet_samples: -1 # jet samples to generate
+  #  model_name: "lhco_flow_matching_both_jet"
+  #  log_epoch_zero: True
+  #  data_type: "val"
+  #  w_dist_config:
+  #    num_eval_samples: 10_000
+  #    num_batches: 40
+  #    calculate_efps: False
+  #  generation_config:
+  #    batch_size: 2048
+  #    ode_solver: "midpoint"
+  #    ode_steps: 100
+  #    verbose: True
+  #  plot_config:
+  #    plot_efps: False
+  #    plottype: ""
+
+  lhco_eval_sr:
+    every_n_epochs: 300 # evaluate every n epochs
+    num_jet_samples: -1 # jet samples to generate
+    model_name: "lhco_flow_matching_both_jet_sr"
+    log_epoch_zero: True
+    data_type: "val"
+    w_dist_config:
+      num_eval_samples: 10_000
+      num_batches: 40
+      calculate_efps: False
+    generation_config:
+      batch_size: 2048
+      ode_solver: "midpoint"
+      ode_steps: 100
+      verbose: True
+    plot_config:
+      plot_efps: False
+      plottype: ""
+
+task_name: "lhco_flow_matching_both_jet_crossattention"
+
+logger:
+  wandb:
+    tags: ${tags}
+    project: "LHCO"
+    group: "lhco_flow_matching"
+    name: ${task_name}
+  comet:
+    project_name: "LHCO"
+    experiment_name: ${task_name}

configs/experiment/lhco/lhco_both_jets_transformer.yaml

@@ -0,0 +1,118 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=fm_tops
+
+defaults:
+  - override /data: lhco.yaml
+  - override /model: fm_droid_transformer.yaml
+  - override /callbacks: /lhco/lhco.yaml
+  - override /trainer: ddp.yaml
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+
+
+tags: ["flow_matching", "LHCO", "cond", "both_jet", "transformer"]
+
+seed: 12345
+
+trainer:
+  min_epochs: 250
+  max_epochs: 5000
+  gradient_clip_val: 0.5
+
+model:
+  num_particles: 279
+  global_cond_dim: 5 # needs to be calculated when using conditioning
+  scheduler:
+    warmup: ${trainer.min_epochs}
+    max_iters: ${trainer.max_epochs}
+  #optimizer:
+  #  _target_: torch.optim.AdamW
+  #  _partial_: true
+  #  lr: 0.0003
+  #  weight_decay: 0.00005
+
+  #scheduler:
+  #  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  #  _partial_: true
+  #  T_max: ${trainer.max_epochs}
+
+data:
+  batch_size: 128
+  normalize: True
+  normalize_sigma: 1
+  conditioning: True
+  relative_coords: True
+  jet_type: "all"
+  num_particles: ${model.num_particles}
+  val_fraction: 0.05
+  test_fraction: 0.35
+  use_all_data: False
+  shuffle_data: False
+  file_suffix_processed_data: "_masssorted"
+  log_pt: True
+  pt_standardization: False
+  multiplicity_conditioning: True
+
+callbacks:
+  ema:
+    decay: 0.999
+    apply_ema_every_n_steps: 1
+    start_step: 0
+    save_ema_weights_in_callback_state: True
+    evaluate_ema_weights_instead: True
+
+  early_stopping: null
+
+  #lhco_eval:
+  #  every_n_epochs: 300 # evaluate every n epochs
+  #  num_jet_samples: -1 # jet samples to generate
+  #  model_name: "lhco_flow_matching_both_jet"
+  #  log_epoch_zero: True
+  #  data_type: "val"
+  #  w_dist_config:
+  #    num_eval_samples: 10_000
+  #    num_batches: 40
+  #    calculate_efps: False
+  #  generation_config:
+  #    batch_size: 2048
+  #    ode_solver: "midpoint"
+  #    ode_steps: 100
+  #    verbose: True
+  #  plot_config:
+  #    plot_efps: False
+  #    plottype: ""
+
+  lhco_eval_sr:
+    every_n_epochs: 300 # evaluate every n epochs
+    num_jet_samples: -1 # jet samples to generate
+    model_name: "lhco_flow_matching_both_jet_sr"
+    log_epoch_zero: True
+    data_type: "val"
+    w_dist_config:
+      num_eval_samples: 10_000
+      num_batches: 40
+      calculate_efps: False
+    generation_config:
+      batch_size: 2048
+      ode_solver: "midpoint"
+      ode_steps: 100
+      verbose: True
+    plot_config:
+      plot_efps: False
+      plottype: ""
+
+task_name: "lhco_flow_matching_both_jet_transformer"
+
+logger:
+  wandb:
+    tags: ${tags}
+    project: "LHCO"
+    group: "lhco_flow_matching"
+    name: ${task_name}
+  comet:
+    project_name: "LHCO"
+    experiment_name: ${task_name}